package zuc

import (
	"reflect"
	"testing"
)

func Test_genKeyword_case1(t *testing.T) {
	s, _ := newZUCState([]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})
	z1 := s.genKeyword()
	if z1 != 0x27bede74 {
		t.Errorf("expected=%x, result=%x\n", 0x27bede74, z1)
	}
	z2 := s.genKeyword()
	if z2 != 0x018082da {
		t.Errorf("expected=%x, result=%x\n", 0x018082da, z2)
	}
}

func Test_genKeyword_case2(t *testing.T) {
	s, _ := newZUCState([]byte{255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255}, []byte{255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255})
	z1 := s.genKeyword()
	if z1 != 0x0657cfa0 {
		t.Errorf("expected=%x, result=%x\n", 0x0657cfa0, z1)
	}
	z2 := s.genKeyword()
	if z2 != 0x7096398b {
		t.Errorf("expected=%x, result=%x\n", 0x7096398b, z2)
	}
}

func Test_genKeyword_case3(t *testing.T) {
	s, _ := newZUCState([]byte{0x3d, 0x4c, 0x4b, 0xe9, 0x6a, 0x82, 0xfd, 0xae, 0xb5, 0x8f, 0x64, 0x1d, 0xb1, 0x7b, 0x45, 0x5b}, []byte{0x84, 0x31, 0x9a, 0xa8, 0xde, 0x69, 0x15, 0xca, 0x1f, 0x6b, 0xda, 0x6b, 0xfb, 0xd8, 0xc7, 0x66})
	z1 := s.genKeyword()
	if z1 != 0x14f1c272 {
		t.Errorf("expected=%x, result=%x\n", 0x14f1c272, z1)
	}
	z2 := s.genKeyword()
	if z2 != 0x3279c419 {
		t.Errorf("expected=%x, result=%x\n", 0x3279c419, z2)
	}
}

var zuc256Tests = []struct {
	key          []byte
	iv           []byte
	expectedKeys []uint32
}{
	{
		[]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		},
		[]uint32{
			0x58d03ad6, 0x2e032ce2, 0xdafc683a, 0x39bdcb03, 0x52a2bc67,
			0xf1b7de74, 0x163ce3a1, 0x01ef5558, 0x9639d75b, 0x95fa681b,
			0x7f090df7, 0x56391ccc, 0x903b7612, 0x744d544c, 0x17bc3fad,
			0x8b163b08, 0x21787c0b, 0x97775bb8, 0x4943c6bb, 0xe8ad8afd,
		},
	},
	{
		[]byte{
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		},
		[]byte{
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		},
		[]uint32{
			0x3356cbae, 0xd1a1c18b, 0x6baa4ffe, 0x343f777c, 0x9e15128f,
			0x251ab65b, 0x949f7b26, 0xef7157f2, 0x96dd2fa9, 0xdf95e3ee,
			0x7a5be02e, 0xc32ba585, 0x505af316, 0xc2f9ded2, 0x7cdbd935,
			0xe441ce11, 0x15fd0a80, 0xbb7aef67, 0x68989416, 0xb8fac8c2,
		},
	},
}

func Test_ZUC256(t *testing.T) {
	for i, test := range zuc256Tests {
		c, err := newZUCState(test.key, test.iv)
		if err != nil {
			t.Error(err)
		}
		words := make([]uint32, 20)
		c.genKeywords(words)
		if !reflect.DeepEqual(words, test.expectedKeys) {
			t.Errorf("#%d: wrong output: got %x, expected %x", i, words, test.expectedKeys)
		}
	}
}