// Unit testing for Carbon C2 handler crypto-related functionality

package carbon

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"testing"
	
	"golang.org/x/crypto/cast5"
)

const (
	rsaPubKeyExponent = 0x11
	rsaPubKeyModulusStr = "c5cbeff7c36cb97d45b9f7fd2c3c95e5fa69fcc01b3ef21888cca8ba9b3dba1273eefd04e0c442650a0ceb0e3dae398c4e0b29b374c97bc211ffac1a10e4f37af5419a66b62c577a43ec253a77dd1c52dad988c25326357ef04ffbba6873d53f8637f898bb3d8468a1f3005410f66d3587ab03f6993366f0a47ec98d35dd426099d18ed66a5367de67be63d9845d04b5ebebb15d7dee123074ce0f49fb5f07e5d0d17b8d0ef75a4b4e80062b0ffeb376ecf70baf2f0d2c02b510bac0b44ff9ac50e01ac0f8140c3bf4d3cc353c96e51d19c3ababb5475583ef75630e8d5c19dfce7d91f5eba26454b88e5ecadc8182fcdb472bd41b90b81a0cb61111e5052cef"
		dummyPlaintext = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."
	rsaPrivateKeyStr = `-----BEGIN RSA PRIVATE KEY-----
MIIEoQIBAAKCAQEAxcvv98NsuX1Fuff9LDyV5fpp/MAbPvIYiMyoups9uhJz7v0E
4MRCZQoM6w49rjmMTgsps3TJe8IR/6waEOTzevVBmma2LFd6Q+wlOnfdHFLa2YjC
UyY1fvBP+7poc9U/hjf4mLs9hGih8wBUEPZtNYerA/aZM2bwpH7JjTXdQmCZ0Y7W
alNn3me+Y9mEXQS16+uxXX3uEjB0zg9J+18H5dDRe40O91pLToAGKw/+s3bs9wuv
Lw0sArUQusC0T/msUOAawPgUDDv008w1PJblHRnDq6u1R1WD73VjDo1cGd/OfZH1
66JkVLiOXsrcgYL820cr1BuQuBoMthER5QUs7wIBEQKCAQAdFnYc6Ah1oXsx78NZ
VDQpWYgOlLi2buV9h4I5j0zXmU1IytsR/r54RT4ikScwNaOxH8JeJ8NG59V4bCHz
bPahJBEtS1cGhVW+sck9TdzAZomYdf51o7ySqt6V9cQRCMWTvO/aOacqD2McNMER
jaal/Vzp/p4PFqrrA5YcS6+Y0bZRa2DUwrhC4w6O6F+2TTuCeJy8QvYZ4FUc+mOh
28c8pAHpvOnPUCI9LD27ksjwvkwzQCQH+8+lIebQuRqmQR/bsphPHJhmAxNiXP2B
dfL/WkdkxM9VIKQQyZpjYHa48nlCTop/uu9vyydVr1gkp9OOmPth9nbjk8AAliEl
bD51AoGBAOJOrCsLobcz2YakqoxLeBbuTjWNnSsC/U5GdG7UMOjW3ZtBFX0TrQMp
GmW3r9UH94tWHVrl7iCWsn2BspARw1xAoTYzIvCiYoR51qiFGRrlncmr6WQE+esb
gRVJHS+BuDNhr7OxXlE5726OZHvOBlMxK5sFLJ47yh7L0oWdti1zAoGBAN+/ohjr
HzIW7KGNAtOgTD2GaVIC5jmScOPCjc9A8Tqlyyk4P8Jh8sW4ny/eRtNGcVt3oJJ5
O4dvgnGtvQige3dtgHJz332A97lWsGp6W7w74uFSiAKZFz0umWchrQVIHS9Y/2E8
GbbvY63wJG+6OqStPn0BljBwyaEZdN4VoiOVAoGAXS90Ebl+0vc7c603KrWp61MR
JRwxqEyGa4ZsLaKqujpbP+2fb7zOxRDswHjP7k6TG0GTneY04D4NQrzvLEOMrYQG
JWBZrmD7Y7myvdxzv8f1rWTnoaeyM6Hp25aTjAg8ydzt/rJyIXI1acIpYCeoQF+K
bQIhblTawWL8VSLSiy8CgYEAqxoSi4afYooAP0223hErPht9ty9kwp0pJqPV2rkw
8JzmpwzldocjD6tMjgRURzXeNuMCUeQ8lL6vC6L594nH08w1DDp9ulOQQm932PQo
CGoH2XtY8u2KPdhXML9mMTclYHE7wtObMYni0E45+xXwnAwCm9QJcFY/1Yvv9R+a
GzUCgYBSb0kAPXlL7ZkwuTxfbvc10/93Ks8LDd5WaAb+gnTDFhqFGjNYNRsSF3S0
9oqfoITt0t4ufZfu4uqtDMFfCCmLA6K2J3asFSFV9A57f4NNtNivgMeoJFsWmLiW
0obQRCbpQ1DY3AcgYPuiI8sTS0bobizCA3MenIWpyMlXT71VvQ==
-----END RSA PRIVATE KEY-----`
)

var (
	// Generated using: echo -n [lorem ipsum text] | openssl enc -e -cast5-cbc -iv "0000000000000000" -K "00000000000000000000000000000000" -provider legacy | xxd -p -i -c 10000
	// also prepended 8 zero bytes for the IV
	dummyCiphertext = []byte{
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xa2, 0xb6, 0x25, 0x85, 0x88, 0xac, 0x5b, 0x64, 0xc1, 0xa5, 0x41, 0x54, 0x39, 0x1f, 0xa0, 0xce, 0x01, 0xea, 0x74, 0x3d, 0x4d, 0xd9, 0x3b, 0xa9, 0x09, 0xba, 0xcb, 0xc2, 0xfb, 0xe3, 0xa0, 0x28, 0xea, 0x77, 0x54, 0x0c, 0x47, 0x5c, 0x9b, 0xad, 0xf9, 0x2e, 0x86, 0x98, 0xb0, 0xc3, 0xc2, 0x41, 0xcf, 0x64, 0xc7, 0x6d, 0xa6, 0x57, 0xcf, 0x1c, 0x2c, 0xbb, 0xbd, 0x8d, 0x36, 0x42, 0xe8, 0xda, 0x05, 0x28, 0x6f, 0xc9, 0x26, 0x82, 0x33, 0xbc, 0x90, 0x20, 0x3b, 0x32, 0x11, 0x1b, 0xa2, 0x2e, 0x7e, 0xd5, 0x16, 0x7d, 0xc2, 0x6b, 0x49, 0x35, 0xe9, 0x44, 0x74, 0xb6, 0x09, 0x70, 0x05, 0x42, 0x24, 0xb4, 0x4d, 0x5d, 0x2e, 0xd6, 0x66, 0x76, 0x3e, 0xba, 0xe9, 0x30, 0x3f, 0x5c, 0xb7, 0x2b, 0x82, 0xdc, 0xfa, 0x16, 0xb7, 0xdc, 0x48, 0x97, 0x3b, 0x34, 0xbc, 0xcf, 0x6c, 0xdf, 0x97, 0xdb, 0x4d, 0xa4, 0x46, 0xbf, 0x83, 0xfb, 0x71, 0xb5, 0x6c, 0x06, 0x7f, 0x48, 0x6c, 0x34, 0x56, 0x2e, 0x14, 0x6c, 0xc5, 0x6e, 0xa3, 0x7b, 0x88, 0x0e, 0x8d, 0x90, 0xa8, 0xba, 0xfa, 0xe5, 0x0b, 0xfa, 0xce, 0x94, 0xd4, 0x9c, 0x45, 0xbb, 0xe0, 0xce, 0x4b, 0x47, 0xc2, 0x0b, 0x06, 0x51, 0xb7, 0x3f, 0x55, 0x9f, 0x02, 0x53, 0x30, 0x9b, 0x75, 0xc4, 0x2e, 0xb9, 0xf7, 0x40, 0xdd, 0xa8, 0xda, 0x87, 0xa9, 0x57, 0x59, 0x61, 0xd1, 0x8a, 0xfe, 0x13, 0x3b, 0xfa, 0xf3, 0xdd, 0xb3, 0x75, 0x91, 0xb8, 0x3a, 0x70, 0x43, 0x6d, 0x0e, 0xe7, 0xce, 0x1f, 0x26, 0x81, 0x46, 0x2f, 0x4f, 0xe3, 0x68, 0xa1, 0xd5, 0xd3, 0x7d, 0xc0, 0x23, 0x5d, 0x9b, 0xfe, 0x22, 0x56, 0x27, 0x89, 0xc8, 0xf5, 0x4b, 0x24, 0x6b, 0x76, 0x49, 0x51, 0xba, 0x63, 0xcf, 0x15, 0x3a, 0x01, 0x62, 0x1d, 0x74, 0x03, 0x9a, 0x56, 0xac, 0x0e, 0xf6, 0x60, 0x15, 0x16, 0xd1, 0x01, 0x90, 0x1a, 0x24, 0xb5, 0x39, 0x08, 0xa3, 0x37, 0x24, 0x3a, 0xd6, 0x38, 0xb0, 0x88, 0x88, 0xe3, 0x69, 0x75, 0x05, 0x15, 0xa0, 0xc1, 0x45, 0x9d, 0x2f, 0xec, 0x88, 0xb6, 0xdc, 0x49, 0x2d, 0xfc, 0x94, 0x26, 0x94, 0x2b, 0x2f, 0xf7, 0x36, 0xc7, 0x06, 0xc9, 0xa7, 0x5f, 0x0a, 0x84, 0x47, 0x09, 0x43, 0x43, 0xa6, 0x69, 0x34, 0x80, 0x08, 0xed, 0x2c, 0x5f, 0xc8, 0x6c, 0x5f, 0xba, 0xef, 0x42, 0xf3, 0xf8, 0x57, 0x43, 0xcb, 0x7e, 0xc4, 0xbf, 0x65, 0x36, 0x60, 0xe1, 0xb7, 0xbe, 0x48, 0x4d, 0x26, 0xb4, 0x08, 0x01, 0x8a, 0x14, 0x92, 0xd6, 0x3e, 0x86, 0x27, 0x6a, 0xa0, 0x90, 0x76, 0x67, 0xe7, 0x19, 0x6a, 0x8e, 0x18, 0x09, 0x1c, 0x10, 0x16, 0x1c, 0xd5, 0xac, 0x43, 0x11, 0xa4, 0x95, 0x94, 0xa4, 0x22, 0xf5, 0xad, 0xda, 0xce, 0x60, 0x44, 0xca, 0xa6, 0xae, 0xb1, 0x67, 0x1d, 0x5e, 0x56, 0x34, 0x45, 0x4e, 0xdd, 0x89, 0xcc, 0x60, 0x79, 0x30, 0x42, 0xa8, 0x65, 0x05, 0x9b, 0xb2, 0x9b, 0x98, 0x41, 0xc4, 0x40, 0x32, 0xb5, 0x7e, 0xd3, 0xfb, 0xf1, 0x6b, 0x15, 0x32, 0x3d, 0xb6, 0x53, 0xd8, 0x15, 0x6c, 0xba, 0xae, 0xfb, 0x87, 0xb6, 0x29, 0x56, 0xd8, 0x3c, 0x2c, 0xf4, 0x46, 0xfe, 0x34, 0xc9, 0x67, 0x5f, 0xab, 0x95, 0x94,
	}
	dummyBytes = []byte{
		0x8a, 0xfe, 0x13, 0x3b, 0xfa, 0xf3, 0xdd, 0xb3, 0x75, 0x91, 0xb8, 0x3a, 0x70, 0x43, 0x6d, 0x0e, 0xe7, 0xce, 0x1f, 0x26, 0x81, 0x46, 0x2f, 0x4f, 0xe3, 0x68, 0xa1, 0xd5, 0xd3, 0x7d, 0xc0, 0x23, 0x5d, 0x9b, 0xfe, 0x22, 0x56, 0x27, 0x89, 0xc8, 0xf5, 0x4b, 0x24, 0x6b, 0x76, 0x49, 0x51, 0xba, 0x63, 0xcf, 0x15, 0x3a, 0x01, 0x62, 0x1d, 0x74, 0x03, 0x9a, 0x56, 0xac, 0x0e, 0xf6, 0x60, 0x15, 0x16, 0xd1, 0x01, 0x90, 0x1a, 0x24, 0xb5, 0x39, 0x08, 0xa3, 0x37, 0x24, 0x3a, 0xd6, 0x38, 0xb0, 0x88, 0x88, 0xe3, 0x69, 0x75, 0x05, 0x15, 0xa0, 0xc1, 0x45, 0x9d, 0x2f, 0xec, 0x88, 0xb6, 0xdc, 0x49, 0x2d, 0xfc, 0x94, 0x26, 0x94, 0x2b, 0x2f, 0xf7, 0x36, 0xc7, 0x06, 0xc9, 0xa7, 0x5f, 0x0a, 0x84, 0x47, 0x09, 0x43, 0x43, 0xa6, 0x69, 0x34, 0x80, 0x08, 0xed, 0x2c, 0x5f, 0xc8, 0x6c, 0x5f, 0xba, 0xef, 0x42, 0xf3, 0xf8, 0x57, 0x43, 0xcb, 0x7e, 0xc4, 0xbf, 0x65, 0x36, 0x60, 0xe1, 0xb7, 0xbe, 0x48, 0x4d, 0x26, 0xb4, 0x08, 0x01, 0x8a, 0x14, 0x92, 0xd6, 0x3e, 0x86, 0x27, 0x6a, 0xa0, 0x90, 0x76, 0x67, 0xe7, 0x19, 0x6a, 0x8e, 0x18, 0x09, 0x1c, 0x10, 0x16, 0x1c, 0xd5, 0xac, 0x43, 0x11, 0xa4, 0x95, 0x94, 0xa4, 0x22, 0xf5, 0x27, 0x89, 0xc8, 0xf5, 0x4b, 0x24, 0x6b, 0x76, 0x49, 0x51, 0xba, 0x63, 0xcf, 0x15, 0x3a, 0x01, 0x62, 0x1d, 0x74, 0x03, 0x9a, 0x56, 0xac, 0x0e, 0xf6, 0x60, 0x15, 0x16, 0xd1, 0x01, 0x90, 0x1a, 0x24, 0xb5, 0x39, 0x08, 0xa3, 0x37, 0x24, 0x3a, 0xd6, 0x38, 0xb0, 0x88, 0x88, 0xe3, 0x69, 0x75, 0x05, 0x15, 0xa0, 0xc1, 0x45, 0x9d, 0x2f, 0xec, 0x88, 0xb6, 0xdc, 0x49, 0x2d, 0xfc, 0x94, 0x26, 0x94, 0x2b, 0x2f, 0xf7, 0x36, 0xc7, 0x06, 0xc9, 0xa7, 0x5f, 0x0a, 0x84, 0x47, 0x09, 0x43, 0x43, 0xa6, 0x69, 0x34, 0x80, 0x08, 0xed, 0x2c, 0x5f, 0xc8, 0x6c, 0x5f, 0xba, 0xef, 0x42, 0xf3, 0xf8, 0x57, 0x43, 0xcb, 0x7e, 0xc4, 0xbf, 0x65, 0x36, 0x60, 0xe1, 0xb7, 0xbe, 0x48, 0x4d, 0x26, 0xb4, 0x08, 0x01, 0x8a, 0x14, 0x92, 0xd6, 0x3e, 0x86, 0x27, 0x6a, 0xa0, 0x90, 0x76, 0x67, 0xe7, 0x19, 0x6a, 0x8e, 0x18, 0x09, 0x1c, 0x10, 0x16, 0x1c, 0xd5, 0xac, 0x43, 0x11, 0xa4, 0x95, 0x94, 0xa4, 0x22, 0xf5, 0xad, 0xda, 0xce, 0x60, 0x44, 0xca, 0xa6, 0xae, 0xb1, 0x67, 0x1d, 0x5e, 0x56, 0x34, 0x45, 0x4e, 0xdd, 0x89, 0xcc, 0x60, 0x79, 0x30, 0x42, 0xa8, 0x65, 0x05, 0x9b, 0xb2, 0x9b, 0x98, 0x41, 0xc4, 0x40,
	}
	keyWant = []byte{
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	}

	// Generated using: echo "0000000000000000"$(echo -n "this is my post data" | openssl enc -e -cast -provider legacy -iv 0000000000000000 -K f2d4560891bd948692c28d2a9391e7d9 | xxd -p -c 10000) | xxd -r -p | base64
	dummyPostResponse = []byte("AAAAAAAAAADeDmUodN6hxfEwWJNIKeMN/SywlhRoLSk=")
)

func rsaOaepSha1Decrypt(ciphertext []byte) ([]byte, error) {
	// Import RSA private key
	keyBlock, _ := pem.Decode([]byte(rsaPrivateKeyStr))
	if keyBlock == nil {
		return nil, errors.New("Failed to decode PEM key block")
	}
	privKey, err := x509.ParsePKCS1PrivateKey(keyBlock.Bytes)
	if err != nil {
		return nil, err
	}
	rng := rand.Reader
	return rsa.DecryptOAEP(sha1.New(), rng, privKey, ciphertext, nil)
}

func decodeTaskResponse(taskResp string, handler *CarbonHttpHandler) ([]byte, error) {
	combinedCiphertexts, err := base64.StdEncoding.DecodeString(taskResp)
	if err != nil {
		return nil, err
	}
	if len(combinedCiphertexts) <= 256 {
		return nil, errors.New("Decoded task response too short to contain both ciphertexts")
	}
	
	// first 256 bytes are RSA ciphertext, the rest is CAST5 ciphertext
	rsaCiphertext := combinedCiphertexts[:256]
	castCiphertext := combinedCiphertexts[256:]
	
	// RSA decrypt and then base64 decode to get the key (should be all 0s)
	encodedKey, err := rsaOaepSha1Decrypt(rsaCiphertext)
	if err != nil {
		return nil, err
	}
	castKey, err := base64.StdEncoding.DecodeString(string(encodedKey))
	if err != nil {
		return nil, err
	}
	if bytes.Compare(castKey, keyWant) != 0 {
		return nil, errors.New(fmt.Sprintf("Decrypted key not equal to expected key. Expected: %v, got: %v", keyWant, castKey))
	}
	
	// decrypt CAST ciphertext
	return handler.cast5CbcDecrypt(castKey, castCiphertext)
}

func TestLoadRsaPublicKey(t *testing.T) {
	handler := carbonHttpHandlerFactory(false, mockGenRandBytes)
	
	// Test correct RSA key loading
	modulusStr := handler.getRsaPubKeyModulus().Text(16)
	if modulusStr != rsaPubKeyModulusStr {
		t.Errorf("Wrong RSA public key modulus. Expected %s, got %s", rsaPubKeyModulusStr, modulusStr)
	}
	exponent := handler.getRsaPubKeyExponent()
	if exponent != rsaPubKeyExponent {
		t.Errorf("Wrong RSA public key exponent. Expected %d, got %d", rsaPubKeyExponent, exponent)
	}
}

func TestCast5EncDec(t *testing.T) {
	handler := carbonHttpHandlerFactory(false, mockGenRandBytes)
	
	// generate all-zeroes CAST128 key
	castKey := make([]byte, cast5.KeySize)
	_, err := handler.genRandBytesFn(castKey)
	if err != nil {
		t.Errorf("Failed to create CAST key for testing: %s", err.Error())
	}

	
	// Encrypt dummy data using CAST128. Will use all-0 IV
	ciphertext, err := handler.cast5CbcEncrypt(castKey, []byte(dummyPlaintext))
	if err != nil {
		t.Errorf("Failed to encrypt dummy data %s", err.Error())
	}
	
	// verify ciphertext
	if bytes.Compare(ciphertext, dummyCiphertext) != 0 {
		t.Errorf("Produced ciphertext did not match expected ciphertext.\nExpected:\n%v\n\nGot:\n%v\n", dummyCiphertext, ciphertext)
	}
	
	// decrypt and verify
	plaintext, err := handler.cast5CbcDecrypt(castKey, ciphertext)
	if bytes.Compare(plaintext, []byte(dummyPlaintext)) != 0 {
		t.Errorf("Produced plaintext did not match expected plaintext.\nExpected:\n%s\n\nGot:\n%v\n", dummyPlaintext, plaintext)
	}
}

func TestEncodedTaskResponse(t *testing.T) {
	handler := carbonHttpHandlerFactory(true, mockGenRandBytes)
	
	resp, err := handler.encodeTaskResponse(dummyBytes)
	if err != nil {
		t.Errorf("Failed to encode task response: %s", err.Error())
	}
	decrypted, err := decodeTaskResponse(resp, handler)
	if err != nil {
		t.Errorf("Failed to decode task response: %s", err.Error())
	} else if bytes.Compare(decrypted, []byte(dummyBytes)) != 0 {
		t.Errorf("Decrypted response did not match expected response.\nExpected:\n%v\n\nGot:\n%v\n", dummyBytes, decrypted)
	}
}

func TestDecodingPostResponse(t *testing.T) {
	handler := carbonHttpHandlerFactory(true, mockGenRandBytes)

	resp, err := handler.decodePostResponse(dummyPostResponse)
	if err != nil {
		t.Errorf("Failed to decode post response: %s", err.Error())
	}

	if (string(resp) != "this is my post data") {
		t.Errorf("Decrypted post response did not match expected response.\nExpected: \n%s\nGot: \n%s", "this is my post data", string(resp))
	}
}
