From ea698259d5f50eda7e58705cc69bf8c74e85536f Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Tue, 23 Mar 2021 22:06:39 +0300
Subject: [PATCH] Add raw version of server handshake test

---
 .../obfuscated2/server_handshake_test.go      | 95 +++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 mtglib/internal/obfuscated2/server_handshake_test.go

diff --git a/mtglib/internal/obfuscated2/server_handshake_test.go b/mtglib/internal/obfuscated2/server_handshake_test.go
new file mode 100644
index 0000000..d0b40c4
--- /dev/null
+++ b/mtglib/internal/obfuscated2/server_handshake_test.go
@@ -0,0 +1,95 @@
+package obfuscated2_test
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"encoding/base64"
+	"testing"
+
+	"github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2"
+	"github.com/9seconds/mtg/v2/testlib"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
+)
+
+type ServerHandshakeTestSuite struct {
+	suite.Suite
+}
+
+func (suite *ServerHandshakeTestSuite) TestOk() {
+	buf := &bytes.Buffer{}
+	connMock := &testlib.NetConnMock{}
+
+	encryptor, decryptor, err := obfuscated2.ServerHandshake(buf)
+	suite.NotNil(encryptor)
+	suite.NotNil(decryptor)
+	suite.NoError(err)
+
+	proxyConn := &obfuscated2.Conn{
+		Conn:      connMock,
+		Encryptor: encryptor,
+		Decryptor: decryptor,
+	}
+
+	serverEncrypted := buf.Bytes()
+
+	decBlock, _ := aes.NewCipher(serverEncrypted[8 : 8+32])
+	serverDecryptor := cipher.NewCTR(decBlock, serverEncrypted[8+32:8+32+16])
+	serverDecrypted := make([]byte, len(serverEncrypted))
+	serverDecryptor.XORKeyStream(serverDecrypted, serverEncrypted)
+
+	suite.Equal("3d3d3Q",
+		base64.RawStdEncoding.EncodeToString(serverDecrypted[8+32+16:8+32+16+4]))
+
+	serverEncryptedReverted := make([]byte, len(serverEncrypted))
+
+	for i := 0; i < 32+16; i++ {
+		serverEncryptedReverted[8+i] = serverEncrypted[8+32+16-1-i]
+	}
+
+	encBlock, _ := aes.NewCipher(serverEncryptedReverted[8 : 8+32])
+	serverEncryptor := cipher.NewCTR(encBlock, serverEncryptedReverted[8+32:8+32+16])
+
+	messageFromTelegram := []byte{1, 2, 3, 4, 5}
+	// messageToTelegram := []byte{10, 11, 13, 14}
+	bufferToRead := make([]byte, 5)
+
+	connMock.
+		On("Read", mock.Anything).
+		Return(5, nil).
+		Once().
+		Run(func(args mock.Arguments) {
+			messageToRead := make([]byte, len(messageFromTelegram))
+			serverEncryptor.XORKeyStream(messageToRead, messageFromTelegram)
+			copy(args.Get(0).([]byte), messageToRead)
+		})
+
+	n, err := proxyConn.Read(bufferToRead)
+	suite.EqualValues(5, n)
+	suite.NoError(err)
+	suite.Equal(messageFromTelegram, bufferToRead)
+
+	messageToTelegram := []byte{10, 11, 12, 13, 14}
+
+	connMock.
+		On("Write", mock.Anything).
+		Return(5, nil).
+		Once().
+		Run(func(args mock.Arguments) {
+			message := make([]byte, len(messageToTelegram))
+			serverDecryptor.XORKeyStream(message, args.Get(0).([]byte))
+			suite.Equal(messageToTelegram, message)
+		})
+
+	n, err = proxyConn.Write(messageToTelegram)
+	suite.EqualValues(5, n)
+	suite.NoError(err)
+
+	connMock.AssertExpectations(suite.T())
+}
+
+func TestServerHandshake(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &ServerHandshakeTestSuite{})
+}