// NewObfuscator creates a new Obfuscator, initializes it with
// a seed message, derives client and server keys, and creates
// RC4 stream ciphers to obfuscate data.
func NewObfuscator(config *ObfuscatorConfig) (obfuscator *Obfuscator, err error) {
	seed, err := MakeSecureRandomBytes(OBFUSCATE_SEED_LENGTH)
	if err != nil {
		return nil, ContextError(err)
	}
	clientToServerKey, err := deriveKey(seed, []byte(config.Keyword), []byte(OBFUSCATE_CLIENT_TO_SERVER_IV))
	if err != nil {
		return nil, ContextError(err)
	}
	serverToClientKey, err := deriveKey(seed, []byte(config.Keyword), []byte(OBFUSCATE_SERVER_TO_CLIENT_IV))
	if err != nil {
		return nil, ContextError(err)
	}
	clientToServerCipher, err := rc4.NewCipher(clientToServerKey)
	if err != nil {
		return nil, ContextError(err)
	}
	serverToClientCipher, err := rc4.NewCipher(serverToClientKey)
	if err != nil {
		return nil, ContextError(err)
	}
	maxPadding := OBFUSCATE_MAX_PADDING
	if config.MaxPadding > 0 {
		maxPadding = config.MaxPadding
	}
	seedMessage, err := makeSeedMessage(maxPadding, seed, clientToServerCipher)
	if err != nil {
		return nil, ContextError(err)
	}
	return &Obfuscator{
		seedMessage:          seedMessage,
		clientToServerCipher: clientToServerCipher,
		serverToClientCipher: serverToClientCipher}, nil
}

func NewEncoderDecoder(Passwd string, Server bool) (encoderFunc, decoderFunc func([]byte) []byte, err error) {
	encoder, erre := rc4.NewCipher([]byte(Passwd))
	if erre != nil {
		return nil, nil, erre
	}
	decoder, errd := rc4.NewCipher([]byte(Passwd))
	if errd != nil {
		return nil, nil, errd
	}

	encoderFunc = func(in []byte) []byte {
		buff := make([]byte, len(in))
		encoder.XORKeyStream(buff, in)
		return buff
	}
	decoderFunc = func(in []byte) []byte {
		buff := make([]byte, len(in))
		decoder.XORKeyStream(buff, in)
		return buff
	}
	if Server {
		return decoderFunc, encoderFunc, nil
	}
	return encoderFunc, decoderFunc, nil
}

func initObfuscatorCiphers(
	seed []byte, config *ObfuscatorConfig) (*rc4.Cipher, *rc4.Cipher, error) {

	clientToServerKey, err := deriveKey(seed, []byte(config.Keyword), []byte(OBFUSCATE_CLIENT_TO_SERVER_IV))
	if err != nil {
		return nil, nil, ContextError(err)
	}

	serverToClientKey, err := deriveKey(seed, []byte(config.Keyword), []byte(OBFUSCATE_SERVER_TO_CLIENT_IV))
	if err != nil {
		return nil, nil, ContextError(err)
	}

	clientToServerCipher, err := rc4.NewCipher(clientToServerKey)
	if err != nil {
		return nil, nil, ContextError(err)
	}

	serverToClientCipher, err := rc4.NewCipher(serverToClientKey)
	if err != nil {
		return nil, nil, ContextError(err)
	}

	return clientToServerCipher, serverToClientCipher, nil
}

func NewChiper(algo, secret string) (*Cipher, error) {
	if algo == "rc4" {
		c, err := rc4.NewCipher(truncateSecretToSize(secret, 32))
		if err != nil {
			return nil, err
		}
		return &Cipher{
			enc: c,
			dec: c,
		}, nil
	} else if algo == "aes" {
		key := truncateSecretToSize(secret, 32)
		c, err := aes.NewCipher(key)
		if err != nil {
			return nil, err
		}
		return &Cipher{
			enc: cipher.NewCFBEncrypter(c, key[:c.BlockSize()]),
			dec: cipher.NewCFBDecrypter(c, key[:c.BlockSize()]),
		}, nil
	}

	cipher, err := rc4.NewCipher([]byte(secret))
	if err != nil {
		return nil, err
	}
	return &Cipher{
		enc: cipher,
		dec: cipher,
	}, nil
}

// 密钥交换
func P_get_seed_req(sess *Session, reader *packet.Packet) []byte {
	tbl, _ := PKT_seed_info(reader)
	// KEY1
	X1, E1 := dh.DHExchange()
	KEY1 := dh.DHKey(X1, big.NewInt(int64(tbl.F_client_send_seed)))

	// KEY2
	X2, E2 := dh.DHExchange()
	KEY2 := dh.DHKey(X2, big.NewInt(int64(tbl.F_client_receive_seed)))

	ret := seed_info{int32(E1.Int64()), int32(E2.Int64())}
	// 服务器加密种子是客户端解密种子
	encoder, err := rc4.NewCipher([]byte(fmt.Sprintf("%v%v", SALT, KEY2)))
	if err != nil {
		log.Critical(err)
		return nil
	}
	decoder, err := rc4.NewCipher([]byte(fmt.Sprintf("%v%v", SALT, KEY1)))
	if err != nil {
		log.Critical(err)
		return nil
	}
	sess.Encoder = encoder
	sess.Decoder = decoder
	sess.Flag |= SESS_KEYEXCG
	return packet.Pack(Code["get_seed_ack"], ret, nil)
}

func main() {
	//明文
	src := []byte("Hello, world!")

	//密钥
	key := []byte("12345")

	cipher, err := rc4.NewCipher(key)
	if err != nil {
		fmt.Println("rc4.NewCipher error:" + err.Error())
	}

	encrypted := make([]byte, len(src))
	cipher.XORKeyStream(encrypted, src)
	fmt.Printf("Encrypting %s : %v -> %v\n", src, []byte(src), encrypted)

	decrypted := make([]byte, len(encrypted))
	cipher, err = rc4.NewCipher(key)
	if err != nil {
		fmt.Println("rc4.NewCipher error:" + err.Error())
	}

	cipher.XORKeyStream(decrypted, encrypted)
	fmt.Printf("Decrypting %v -> %v : %s\n", encrypted, decrypted, decrypted)

}

func Obfs3fHandshake(wire net.Conn, is_server bool) (io.ReadWriteCloser, error) {
	var their_public dh_public_key
	var our_keypair = dh_gen_key(1536)
	var secret []byte
	var write_rc4, read_rc4 cipher.Stream
	their_public = make([]byte, 1536/8)
	if is_server {
		_, err := io.ReadFull(wire, their_public)
		if err != nil {
			return nil, err
		}
		_, err = wire.Write(our_keypair.Public)
		if err != nil {
			return nil, err
		}
		secret = dh_gen_secret(our_keypair.Private, their_public)
		write_rc4, _ = rc4.NewCipher(KeyedHash(secret, []byte("obfs3f_downstr")))
		read_rc4, _ = rc4.NewCipher(KeyedHash(secret, []byte("obfs3f_upstr")))
	} else {
		_, err := wire.Write(our_keypair.Public)
		if err != nil {
			return nil, err
		}
		_, err = io.ReadFull(wire, their_public)
		if err != nil {
			return nil, err
		}
		secret = dh_gen_secret(our_keypair.Private, their_public)
		write_rc4, _ = rc4.NewCipher(KeyedHash(secret, []byte("obfs3f_upstr")))
		read_rc4, _ = rc4.NewCipher(KeyedHash(secret, []byte("obfs3f_downstr")))
	}
	read_rc4.XORKeyStream(make([]byte, 8192), make([]byte, 8192))
	write_rc4.XORKeyStream(make([]byte, 8192), make([]byte, 8192))

	toret := &Obfs3f{read_rc4, write_rc4, wire}
	thing := make(chan bool)
	go func() {
		randlen := make([]byte, 2)
		rand.Read(randlen)
		rlint := int(randlen[0])*256 + int(randlen[1])
		xaxa := make([]byte, rlint)
		toret.Write(randlen)
		toret.Write(xaxa)
		thing <- true
	}()

	randlen := make([]byte, 2)
	_, err := io.ReadFull(toret, randlen)
	if err != nil {
		return nil, err
	}
	_, err = io.ReadFull(toret, make([]byte, int(randlen[0])*256+int(randlen[1])))
	if err != nil {
		return nil, err
	}
	<-thing
	return io.ReadWriteCloser(toret), nil
}

func NewRC4Conn(conn net.Conn, key []byte, iv []byte) (sc net.Conn, err error) {
	in, err := rc4.NewCipher(key)
	if err != nil {
		return
	}
	out, err := rc4.NewCipher(key)
	if err != nil {
		return
	}
	return CryptConn{conn.(*net.TCPConn), in, out}, nil
}

func (c *Connection) SetCipher(key []uint8) (err error) {
	if key == nil {
		c.icipher = nil
		c.ocipher = nil
	} else {
		c.icipher, err = rc4.NewCipher(key)
		c.ocipher, err = rc4.NewCipher(key)
	}

	return
}

func TestAgent(t *testing.T) {
	host := DEFAULT_AGENT_HOST
	if env := os.Getenv("AGENT_HOST"); env != "" {
		host = env
	}
	addr, err := net.ResolveTCPAddr("tcp", host)
	if err != nil {
		log.Println(err)
		os.Exit(-1)
	}
	conn, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		log.Println(err)
		os.Exit(-1)
	}
	defer conn.Close()

	//get_seed_req
	S1, M1 := dh.DHExchange()
	S2, M2 := dh.DHExchange()
	p2 := seed_info{
		int32(M1.Int64()),
		int32(M2.Int64()),
	}
	rst := send_proto(conn, Code["get_seed_req"], p2)
	r1, _ := PKT_seed_info(rst)

	K1 := db.DHKey(S1, big.NewInt(int64(r1.F_client_send_seed)))
	K2 := db.DHKey(S2, big.NewInt(int64(r1.F_client_receive_seed)))
	encoder, err = rc4.NewCipher([]byte(fmt.Sprintf("%v%v", SALT, K1)))
	if err != nil {
		log.Println(err)
		return
	}
	decoder, err = rc4.NewCipher([]byte(fmt.Sprintf("%v%v", SALT, K2)))
	if err != nil {
		log.Println(err)
		return
	}

	KEY_EXCHANGE = true

	// heart_beat_req
	p1 := auto_id{F_id: rand.Int31()}
	log.Println("seed: %#v", p1)
	rst = send_proto(conn, Code["heart_beat_req"], p1)
	r2, _ := PKT_auto_id(rst)
	log.Printf("response: %#v", r2)

}

// LLObfsServerHandshake negotiates low-level obfuscation (content hiding but no
// volume hiding) on a network connection, acting as the server. The master
// secret must be provided.
func LLObfsServerHandshake(secret []byte, transport io.ReadWriteCloser) (io.ReadWriteCloser, error) {
	// Client needs to send proof that they actually have our secret
	proof := make([]byte, 64)
	_, err := io.ReadFull(transport, proof)
	if err != nil {
		return nil, err
	}
	kilog.FineDebug("llobfs: server obtained proof")
	// We need to verify proof
	nonce := proof[:32]
	hash := proof[32:]
	if subtle.ConstantTimeCompare(natrium.SecureHash(secret, nonce), hash) != 1 {
		return nil, errors.New("Client did not give the right proof")
	}
	// Generate our ephemeral keys
	our_keys := UDHGenerateKeys()

	// Send our public key
	_, err = transport.Write(our_keys.Public)
	if err != nil {
		return nil, err
	}
	kilog.FineDebug("llobfs: server sent public key")

	// Read their public key
	their_public := make([]byte, 1536/8)
	_, err = io.ReadFull(transport, their_public)
	if err != nil {
		return nil, err
	}
	kilog.FineDebug("llobfs: server read public key")
	// Compute shared secret
	shared_secret := UDHSecret(our_keys.Private, their_public)
	// Read and write keys
	read_key := natrium.SecureHash(shared_secret, []byte("llobfs-upstream-key"))
	write_key := natrium.SecureHash(shared_secret, []byte("llobfs-downstream-key"))
	// Create struct
	toret := new(llobfsContext)

	toret.read_chug, _ = rc4.NewCipher(read_key)
	toret.write_chug, _ = rc4.NewCipher(write_key)

	dummy := make([]byte, 1536)
	toret.read_chug.XORKeyStream(dummy, dummy)
	toret.write_chug.XORKeyStream(dummy, dummy)

	toret.underlying = transport
	return toret, nil
}

func (s *Stream) initRC4(encKey, decKey string, S *big.Int, sKey []byte) error {
	cipherEnc, err := rc4.NewCipher(rc4Key(encKey, S, sKey))
	if err != nil {
		return err
	}
	cipherDec, err := rc4.NewCipher(rc4Key(decKey, S, sKey))
	if err != nil {
		return err
	}
	cipherEnc.XORKeyStream(discard, discard)
	cipherDec.XORKeyStream(discard, discard)
	s.w = &cipher.StreamWriter{S: cipherEnc, W: s.raw}
	s.r = &cipher.StreamReader{S: cipherDec, R: s.raw}
	return nil
}

func read_stego(img io.Reader, secret string, out io.Writer) {

	// see if encryption is used or not
	var encrypt bool = false
	var crypt *rc4.Cipher

	if secret != "" {
		encrypt = true
		var err error
		crypt, err = rc4.NewCipher([]byte(secret))
		if err != nil {
			panic(err)
		}

		// when the function is finished, zero the keyspace
		defer crypt.Reset()
	}

	// get a buffer to read data from the image
	buf := make([]byte, 128)

	// create the stego img reader to read the data
	stegoReader, err := stegoimg.NewStegoImgReader(img)
	if err != nil {
		panic(err)
	}

	// read and decrypt the data
	for {

		// read the new data block
		n, err := stegoReader.Read(buf)
		if err != nil && err != io.EOF {
			panic(err)
		}

		// break if no data was read
		if n == 0 {
			break
		}

		// optionally decrypt
		if encrypt {
			crypt.XORKeyStream(buf[:n], buf[:n])
		}

		// write to the stego writer
		_, writeErr := out.Write(buf[:n])
		if writeErr != nil {
			panic(writeErr)
		}

		// break if that's the end of the data
		if err == io.EOF {
			break
		}
	}

	return
}

// Check the password of the player.
// Return false if connection shall be disonnected
func (up *user) CmdPassword_WLwWLuWLqBlWLc(encrPass []byte) bool {
	// The password is given by the client as an encrypted byte vector.
	// fmt.Printf("CmdPassword: New player encr passwd%v\n", encrPass)
	// Decrypt the password using the full license key.
	cipher, err := rc4.NewCipher(xorVector([]byte(up.License), up.challenge))
	if err != nil {
		log.Printf("CmdPassword: NewCipher1 returned %v\n", err)
		return false
	}
	passw := make([]byte, len(encrPass))
	cipher.XORKeyStream(passw, encrPass)
	// fmt.Printf("CmdPassword: Decrypted password is %#v\n", string(passw))
	if !license.VerifyPassword(string(passw), up.Password, encryptionSalt) {
		// fmt.Println("CmdPassword: stored password doesn't match the given")
		// CmdLogin_WLwWLuWLqBlWLc(up.Name, index)
		if *verboseFlag > 0 {
			log.Println("Terminate because of bad password")
		}
		return false
	}
	// Save player logon time
	up.Lastseen = time.Now()
	db := ephenationdb.New()
	err = db.C("avatars").UpdateId(up.Id, bson.M{"$set": bson.M{"lastseen": up.Lastseen}})
	if err != nil {
		log.Println("Update lastseen", err)
	}
	up.loginAck_WLuWLqBlWLa()
	return true
}

// Encipher a message
func EncipherMsg(plaintext []byte, pubkey *CipherMsgPubKey) ([]byte, os.Error) {
	msg := &U_CipherMsg{
		Text: make([]byte, len(plaintext)),
	}
	n := copy(msg.Text, plaintext)
	if n != len(plaintext) {
		panic("crypto, copy text")
	}
	urand := NewTimedRand()
	seed := make([]byte, cipherMsgSeedLen)
	n, _ = urand.Read(seed)
	if n != len(seed) {
		panic("crypto,gen seed")
	}
	cseed, err := EncryptShortMsg(pubkey.rsa, seed, []byte(""))
	if err != nil {
		return nil, err
	}
	msg.Seed = cseed
	rc, err := rc4.NewCipher(seed)
	if err != nil {
		panic("rc4tube")
	}
	rc.XORKeyStream(msg.Text)
	var w bytes.Buffer
	enc := gob.NewEncoder(&w)
	err = enc.Encode(msg)
	if err != nil {
		return nil, err
	}
	return w.Bytes(), nil
}

func (ev *EncryptEventV2) Decode(buffer *bytes.Buffer) (err error) {
	ev.EncryptType, err = DecodeUInt32Value(buffer)
	if err != nil {
		return err
	}
	length, err := DecodeUInt32Value(buffer)
	if err != nil {
		return
	}
	switch ev.EncryptType {
	case ENCRYPTER_NONE:
		err, ev.Ev = DecodeEvent(buffer)
		return err
	case ENCRYPTER_SE1:
		newbuf := util.SimpleDecrypt(bytes.NewBuffer(buffer.Next(int(length))))
		err, ev.Ev = DecodeEvent(newbuf)
		newbuf.Reset()
	case ENCRYPTER_RC4:
		src := buffer.Next(int(length))
		dst := make([]byte, int(length))
		cipher, _ := rc4.NewCipher([]byte(rc4Key))
		cipher.XORKeyStream(dst, src)
		err, ev.Ev = DecodeEvent(bytes.NewBuffer(dst))
	default:
		return errors.New("Not supported encrypt type:" + strconv.Itoa(int(ev.EncryptType)))
	}
	return err
}

func newRc4Stream(seed string) (*rc4Stream, error) {
	cipher, err := rc4.NewCipher([]byte(seed))
	if err != nil {
		return nil, err
	}
	return &rc4Stream{cipher}, nil
}

func NewTimedRC4Rand() TimedRC4Rand {
	ciph, err := rc4.NewCipher(bytes.Int64ToBytes(time.Nanoseconds()))
	if err != nil {
		panic("rc4 gen error")
	}
	return TimedRC4Rand{ciph}
}

func (o arc4Verifier) validate(count string, offset uint64, key, plaintext, expectedCiphertext []byte) {
	if offset%16 != 0 || len(plaintext) != 16 || len(expectedCiphertext) != 16 {
		panic(fmt.Errorf("Unexpected input value encountered: offset=%v; len(plaintext)=%v; len(expectedCiphertext)=%v",
			offset,
			len(plaintext),
			len(expectedCiphertext)))
	}
	stream, err := rc4.NewCipher(key)
	if err != nil {
		panic(err)
	}

	var currentOffset uint64 = 0
	ciphertext := make([]byte, len(plaintext))
	for currentOffset <= offset {
		stream.XORKeyStream(ciphertext, plaintext)
		currentOffset += uint64(len(plaintext))
	}
	if !bytes.Equal(ciphertext, expectedCiphertext) {
		panic(fmt.Errorf("vector mismatch @ COUNT = %s:\n  %s != %s\n",
			count,
			hex.EncodeToString(expectedCiphertext),
			hex.EncodeToString(ciphertext)))
	}
}

func (p *protectType) rc4(n uint32, buf *[]byte) {
	if p.rc4cipher == nil || p.rc4n != n {
		p.rc4cipher, _ = rc4.NewCipher(p.objectKey(n))
		p.rc4n = n
	}
	p.rc4cipher.XORKeyStream(*buf, *buf)
}