Merge pull request #285 from Edznux/edz-refactor-1

Refactor of util_agent_message_actions_staging_rsa.go

.gitignore

@@ -25,6 +25,7 @@ venv.bak/
 # ssl certs
 ssl/
 # Mythic files
+mythic-docker/src/Mythic
 mythic_access.*
 mythic_sync/
 postgres-docker/database/

mythic-docker/src/crypto/rsa.go

@@ -6,65 +6,58 @@ import (
 	"crypto/sha1"
 	"crypto/x509"
 	"encoding/pem"
-	"errors"
 	"fmt"
-
-	"github.com/its-a-feature/Mythic/logging"
 )
 
 func GenerateRSAKeyPair() ([]byte, *rsa.PrivateKey, error) {
-	if serverPrivKey, err := rsa.GenerateKey(rand.Reader, 4096); err != nil {
-		logging.LogError(err, "Failed to generate a new RSA keypair")
-		errorString := fmt.Sprintf("Failed to generate a new RSA keypair: %s", err.Error())
-		return nil, nil, errors.New(errorString)
-	} else {
-		serverPubKey := &serverPrivKey.PublicKey
-		pubASN1 := x509.MarshalPKCS1PublicKey(serverPubKey)
-		pubPem := pem.EncodeToMemory(
-			&pem.Block{
-				Type:  "RSA PUBLIC KEY",
-				Bytes: pubASN1,
-			},
-		)
-		return pubPem, serverPrivKey, nil
+	serverPrivKey, err := rsa.GenerateKey(rand.Reader, 4096)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to generate a new RSA keypair: %v", err)
 	}
+
+	serverPubKey := &serverPrivKey.PublicKey
+	pubASN1 := x509.MarshalPKCS1PublicKey(serverPubKey)
+	pubPem := pem.EncodeToMemory(
+		&pem.Block{
+			Type:  "RSA PUBLIC KEY",
+			Bytes: pubASN1,
+		},
+	)
+	return pubPem, serverPrivKey, nil
 }
 
 func RsaDecryptCipherBytes(encryptedData []byte, privateKey *rsa.PrivateKey) ([]byte, error) {
 	hash := sha1.New()
-	if decryptedData, err := rsa.DecryptOAEP(hash, rand.Reader, privateKey, encryptedData, nil); err != nil {
-		logging.LogError(err, "Failed to decrypt with RSA private key")
-		stringErr := errors.New(fmt.Sprintf("Failed to decrypt with RSA private key: %s", err.Error()))
-		return nil, stringErr
-	} else {
-		return decryptedData, nil
+	decryptedData, err := rsa.DecryptOAEP(hash, rand.Reader, privateKey, encryptedData, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decrypt with RSA private key: %v", err)
 	}
+
+	return decryptedData, nil
 }
 
 func RsaEncryptBytes(plainBytes []byte, publicKey []byte) ([]byte, error) {
 	hash := sha1.New()
 	//logging.LogInfo("about to parse public key in RsaEncryptBytes", "public key", publicKey)
-	if pkcs1RSAPublicKey, _ := pem.Decode(publicKey); pkcs1RSAPublicKey == nil {
-		logging.LogError(nil, "Failed to find PEM encoded public key")
-		return nil, errors.New("Failed to find PEM encoded public key")
-	} else {
-		var pubKey *rsa.PublicKey
-		var err error
-		if pubKey, err = x509.ParsePKCS1PublicKey(pkcs1RSAPublicKey.Bytes); err != nil {
-			if pubAny, err := x509.ParsePKIXPublicKey(pkcs1RSAPublicKey.Bytes); err != nil {
-				logging.LogError(err, "Failed to parse public key to encrypt with RSA")
-				errorString := fmt.Sprintf("Failed to parse public key to encrypt with RSA: %s", err.Error())
-				return nil, errors.New(errorString)
-			} else {
-				pubKey = pubAny.(*rsa.PublicKey)
-			}
-		}
-		if encryptedData, err := rsa.EncryptOAEP(hash, rand.Reader, pubKey, plainBytes, nil); err != nil {
-			logging.LogError(err, "Failed to encrypt with RSA key")
-			errorString := fmt.Sprintf("Failed to encrypt with RSA key: %s", err.Error())
-			return nil, errors.New(errorString)
-		} else {
-			return encryptedData, nil
+	pkcs1RSAPublicKey, _ := pem.Decode(publicKey)
+	if pkcs1RSAPublicKey == nil {
+		return nil, fmt.Errorf("failed to find PEM encoded public key")
+	}
+
+	pubKey, err := x509.ParsePKCS1PublicKey(pkcs1RSAPublicKey.Bytes)
+	if err != nil {
+		// Fallback to parsing PKIX instead of PKCS1
+		pubAny, err := x509.ParsePKIXPublicKey(pkcs1RSAPublicKey.Bytes)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse public key to encrypt with RSA: %v", err)
 		}
+		pubKey = pubAny.(*rsa.PublicKey)
+	}
+
+	encryptedData, err := rsa.EncryptOAEP(hash, rand.Reader, pubKey, plainBytes, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to encrypt with RSA key: %v", err)
 	}
+
+	return encryptedData, nil
 }

mythic-docker/src/crypto/rsa_test.go

@@ -0,0 +1,59 @@
+package mythicCrypto
+
+import (
+	"os"
+	"reflect"
+	"testing"
+)
+
+func TestRsaEncryptBytes(t *testing.T) {
+	t.Parallel()
+	type args struct {
+		plainBytes []byte
+		publicKey  string
+	}
+
+	tests := []struct {
+		name    string
+		args    args
+		wantLen int
+		wantErr bool
+	}{
+		{
+			name: "test RSA 4096 key success",
+			args: args{
+				plainBytes: []byte("test"),
+				publicKey:  "./testdata/test_key.pub",
+			},
+			wantLen: 512,
+			wantErr: false,
+		},
+		{
+			name: "test key failure",
+			args: args{
+				plainBytes: []byte("test"),
+				publicKey:  "./testdata/invalid.pub",
+			},
+			wantLen: 0,
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		// shadowing tt to avoid races
+		// see https://gist.github.com/posener/92a55c4cd441fc5e5e85f27bca008721
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			pubPem, _ := os.ReadFile(tt.args.publicKey)
+			got, err := RsaEncryptBytes(tt.args.plainBytes, pubPem)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("RsaEncryptBytes() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(len(got), tt.wantLen) {
+				t.Errorf("RsaEncryptBytes() = \n%v\n, want \n%v\n", len(got), tt.wantLen)
+			}
+		})
+	}
+}

mythic-docker/src/crypto/testdata/invalid.pub

@@ -0,0 +1 @@
+invalid key

mythic-docker/src/crypto/testdata/test_key

@@ -0,0 +1,51 @@
+-----BEGIN RSA PRIVATE KEY-----
+MIIJKAIBAAKCAgEA0c4/6tzhoWq1SvmTD2+IpZw80H8ofOGHYEa+SZtzeLAJQWSP
++DwOnPNt3M0/g8FOQv7cD0htWejeW+T2DH3GkPwFy9d80npF2mPDvPzI2ydb3vbg
+hqJ8P1u6Wb4J1RJLgwGDiCVEeTTf12/gNhyo7rr7lr7Pv7nfAv9Jua+HzPQEw/ZW
+0NNeFLUaA55/Thf1JcT8+GGlhlZaYUKf7wBN/kuLKgLtW2CMhFb6NqpZRkdk81wl
+dJMb/PdcZ+ktIlGX2ioFgEBS7Jy+7YJaCPYlb+CtqWACtj9aSSkEFb7V7U7Qdlel
+1LJV+eAVTSH+lyNBJ+uO5FXvs/7FEJTZiAp5tW0x+PAI9kZiD9Aic9kKCxjForo8
+ts4uhWRuX84RCAoWHeLsdmjGadGhqYnbfNmBw24B9J3vb2/xbSmSS6NmqF/TTZ8D
+XyVUftU/YByyu0Szq7wUTON0zsB8nIE6VN+OBwY5p0FC3hDjYxKUROTnL6ENcG1Y
+KbIa6pPeExV19DoDEa6EQZ2XXWOAvNlhwmymXXicdG0QkELFx6oM8cwguu755ubB
+alNMRA5fIxFlnmzGnF/+P6Af1cvpio9Nx2VVq3rS0/ld4xC/GBX+YD6HaZHRwRLB
+kpoKQx19419Ee8pgd/Fq93jt5iLwH+0wyP4KukyDb/+KmKfeEETJngVH0ZkCAwEA
+AQKCAgAQI9FVtmSKpAB6mjC6niJJWDey8+dH1RsymDqdA78KNKOfotZy+p7GqvDj
+jKwTTRCB1aSU7eXQpfoOIHLfsuDOyD6DaAo/kF21j1lwXVNvzYszH7OV0+PiRWKU
+cwrDWA8UTs5hL7nA58UFCCOb0hBaXirK1da6SXibHMu9g7/7D7Yao6WZ1rrM9cIg
+r2WceroWi90ImA1xZEz7YLJYQuIdGw9jfdpn3kltCFpjEG+B+S3OTWNJwKPlbduP
+WFgrJNL9WuB1tW0Tq/rBn3ag7Buk68g8SIko3z6JNIUujZHcOcUOZ493qXj0+iy1
+VHVHB/SkFypUzTodkw6yp/rEjQa//x6QjAqCmW8bac8dlSouWRcEp1sLgYquouB1
+Gp0Me6QXzHR7cukjVdKbB4cGdMAxuMaZTT+tDUKLP0H+Lj3Fts9UISlDqbbdfPLC
+t43iDfGgTttjfF+0CSLCT9xWEctCdB/P8gG9+KK8kFyNjjiljIvxWD4EF3nficz6
+Y2tGk5HbtsdhJ8YJTPDUwD5d7FaIf8R0uWLuS/GeffJAkxXYvxfCIQmMd+HPAHBM
+2pLpyA20RygviTGwiwIVOoK1MaPN/JT05eN37G7BhWAIOzTLHfFZz14oP1FNBZ7O
+VrEIOqROIREfwtritdmk2yj9K8gWxr63E5jbrP3/8IGJV+tAUQKCAQEA+FDzmSe0
++18n+k1UpOKPviPmxu5m9+IT8fbbeGz+S+sVqhKLQoLD9O1ntxkC0JJkggJHg1di
+9T6hx93SBc6JL0Ia+vmhmZiCp0eEVoyeKGcVw0zH1EEoT2iYBZmPpofG2lOpZb4i
+89cvl4Hj3eQI3vSXhdr/eiE1rAVjdrU9qyIBYXKu1UyRQ9aCW7S9woUx7k1Rsgi9
+YspWan03SRxnTLrY779E71fSpaco08UiKRBWCL5J/wcpySzblt9Ls/PClvaB8//y
+IzZQBpjGX9c/ZMp2XAzty/CADPfaxTYgdORM4I1IvSv1vETNfYJLDHGagIYhb65D
+G3G0L0TcUxAskQKCAQEA2Ew8J4iVniyOnHWFJngT0w6QGrLwvvNWytvnkztx1T4V
+QBd/SpDKntqjulIBgEKnPJDlMFPnL8cswdVOQkrvsJMUZAam5Nc14NzrSh7L2A06
+dzRurBzs8qJXsCfq7LPAGrjg+xLJrhhYhI4LgSjTvX7r5fQhvt+WXfJ75WAqrLBQ
+1p6asITH/8un1GQgFIq7S2X6wgErQLeZ0OYk4r+vMdM+PTLuF1ZAljtEL2iMg9Gt
+e51UQKPLxKfKD8x3xzD9+Yn/rkcrs4S7/m8/6yPYKJmMn8tHOu006uTF/QvuCilb
+E3FWw0phcddd/39Ala81mfYrRmIHky4xV4yfPvx4iQKCAQAzYMyGS+jueenTqFxz
+IU2MkfEWCA3WDDkEP4d5i1OycZmx7tRRlqzk6JolEE+8yA1zuPUC//nBtIvUxtjF
+ys4nsQ3UEAKXvS6LXgjLv5yZz3p5RlGyYSjuBT4vKm7GjiFe2yCZpJWmzkdSmdWD
++8K6HvGbCI0DwAtS5GqIDUgNOKQAfbIfABCucJvYTbVVoimKnQTiSVymLCdlMTNX
+NUFKZv6r3G9u9kTncYbJCmjfBjpG5Nh6pAjJTzbAKMOJIE6K3cZVrgdzsbEtLD8f
+ZXIAcMO8mBA7ui9Ef6QMWf6tMO+XJqR+P7JAmhvKdEMC+B53qUkRyoSoEQavIRhV
+N2uBAoIBAFLX2A5YPN3pOPHp/QZ/5S/oGv9u29B3CR4HXcnBcdb3wVHb7hAUJtby
+7NS3BgYnAUCsSmvZJungwl63IRM4+lbJ7nxlI9TwLJ2kX6Xy56YnYuY3OdBH/+Tq
+kuVQVKU2L3TMrLbdOkuo3XZfpT5h8b9ZdmaLu8UMg94Vuqhezdl4am3ZL2w+Xw+0
+4+HwO21CuXumYoWdxwAxkgM1spj6S9KskuEDubdMfM5Ngs8Znv/59hUrbBKZ6bi3
+fyfP62xqckv6M+h/L9jBFPPdjkC0aN0b+oVVaUHDiooQ91f1EsXnkw9+UHnZ6jqc
+/06VYZozEjeW1npNn5MalSFwJaoNLukCggEBANJOP+PoTVIRaZ/0xAs7OBuilzTq
+dWo4hUbpwvlk0x2PuM6sqUSXINR+fWgnLJ1Asq/beEfQSCxpSZLnVcrF7v/8AcO7
+cvCTEJz1drw/XF3iaEoN4Ksai1EfivAGWpdvVMjZJI+ZuHqATzhoyZaKrdG6KwzV
+MeVfaCCOtOZ4dUXhCSFUm7MmPV9nuSM683Q+D9VnbUzXfXx4nbZ5fl4o5gPhl2PI
+MnNvCS8JBpoBxK8ugTn7ecETrP1MuEFioc77f1AA4sL8ztrt61qTcRaS5i3HRQAP
+juxRj4botgqyGbewvPXP8FVzpt5JwBAXx5J7+f12Ezkfr4fOuCfa+MFJ+HY=
+-----END RSA PRIVATE KEY-----

mythic-docker/src/crypto/testdata/test_key.pub

@@ -0,0 +1,13 @@
+-----BEGIN RSA PUBLIC KEY-----
+MIICCgKCAgEA0c4/6tzhoWq1SvmTD2+IpZw80H8ofOGHYEa+SZtzeLAJQWSP+DwO
+nPNt3M0/g8FOQv7cD0htWejeW+T2DH3GkPwFy9d80npF2mPDvPzI2ydb3vbghqJ8
+P1u6Wb4J1RJLgwGDiCVEeTTf12/gNhyo7rr7lr7Pv7nfAv9Jua+HzPQEw/ZW0NNe
+FLUaA55/Thf1JcT8+GGlhlZaYUKf7wBN/kuLKgLtW2CMhFb6NqpZRkdk81wldJMb
+/PdcZ+ktIlGX2ioFgEBS7Jy+7YJaCPYlb+CtqWACtj9aSSkEFb7V7U7Qdlel1LJV
++eAVTSH+lyNBJ+uO5FXvs/7FEJTZiAp5tW0x+PAI9kZiD9Aic9kKCxjForo8ts4u
+hWRuX84RCAoWHeLsdmjGadGhqYnbfNmBw24B9J3vb2/xbSmSS6NmqF/TTZ8DXyVU
+ftU/YByyu0Szq7wUTON0zsB8nIE6VN+OBwY5p0FC3hDjYxKUROTnL6ENcG1YKbIa
+6pPeExV19DoDEa6EQZ2XXWOAvNlhwmymXXicdG0QkELFx6oM8cwguu755ubBalNM
+RA5fIxFlnmzGnF/+P6Af1cvpio9Nx2VVq3rS0/ld4xC/GBX+YD6HaZHRwRLBkpoK
+Qx19419Ee8pgd/Fq93jt5iLwH+0wyP4KukyDb/+KmKfeEETJngVH0ZkCAwEAAQ==
+-----END RSA PUBLIC KEY-----

mythic-docker/src/rabbitmq/util_agent_message_actions_staging_rsa.go

@@ -2,15 +2,13 @@ package rabbitmq
 
 import (
 	"encoding/base64"
-	"errors"
 	"fmt"
 	"strings"
 
 	"github.com/google/uuid"
 	mythicCrypto "github.com/its-a-feature/Mythic/crypto"
 	"github.com/its-a-feature/Mythic/database"
 	databaseStructs "github.com/its-a-feature/Mythic/database/structs"
-	"github.com/its-a-feature/Mythic/logging"
 	"github.com/mitchellh/mapstructure"
 )
 
@@ -20,6 +18,11 @@ type agentMessageStagingRSA struct {
 	Other     map[string]interface{} `json:"-" mapstructure:",remain"`             // capture any 'other' keys that were passed in so we can reply back with them
 }
 
+const (
+	insertQuery = `INSERT INTO staginginfo (session_id, enc_key, dec_key, crypto_type, payload_id, staging_uuid)
+		VALUES (:session_id, :enc_key, :dec_key, :crypto_type, :payload_id, :staging_uuid)`
+)
+
 func handleAgentMessageStagingRSA(incoming *map[string]interface{}, uUIDInfo *cachedUUIDInfo) (map[string]interface{}, error) {
 	// got message:
 	/*
@@ -32,49 +35,51 @@ func handleAgentMessageStagingRSA(incoming *map[string]interface{}, uUIDInfo *ca
 	agentMessage := agentMessageStagingRSA{}
 	stagingDatabaseMessage := databaseStructs.Staginginfo{}
 	if err := mapstructure.Decode(incoming, &agentMessage); err != nil {
-		logging.LogError(err, "Failed to decode agent message into struct")
-		return nil, errors.New(fmt.Sprintf("Failed to decode agent message into struct: %s", err.Error()))
-	} else if newKey, err := mythicCrypto.GenerateKeysForPayload("aes256_hmac"); err != nil {
-		logging.LogError(err, "Failed to generate new AES key for staging rsa")
-		errorString := fmt.Sprintf("Failed to generate new AES key for staging rsa: %s", err.Error())
-		return nil, errors.New(errorString)
-	} else {
-		publicKeyToUse := agentMessage.PublicKey
-		if !strings.HasPrefix(agentMessage.PublicKey, "LS0t") {
-			publicKeyToUse = "-----BEGIN PUBLIC KEY-----\n" + agentMessage.PublicKey + "\n-----END PUBLIC KEY-----"
-		} else if decodedBytes, err := base64.StdEncoding.DecodeString(publicKeyToUse); err != nil {
-			logging.LogError(err, "Failed to base64 provided public key")
-			return nil, err
-		} else {
-			publicKeyToUse = string(decodedBytes)
-		}
-		if encryptedNewKey, err := mythicCrypto.RsaEncryptBytes(*newKey.EncKey, []byte(publicKeyToUse)); err != nil {
-			logging.LogError(err, "Failed to encrypt new encryption key with RSA")
-			return nil, errors.New(fmt.Sprintf("Failed to encrypt new encryption key with RSA: %s", err.Error()))
-		} else if tempUUID, err := uuid.NewRandom(); err != nil {
-			logging.LogError(err, "Failed to generate a new random UUID for staging")
-			return nil, errors.New(fmt.Sprintf("Failed to generate a new random UUID for staging: %s", err.Error()))
-		} else {
-			stagingDatabaseMessage.CryptoType = "aes256_hmac"
-			stagingDatabaseMessage.EncKey = newKey.EncKey
-			stagingDatabaseMessage.DecKey = newKey.DecKey
-			stagingDatabaseMessage.SessionID = agentMessage.SessionID
-			stagingDatabaseMessage.StagingUuID = tempUUID.String()
-			stagingDatabaseMessage.PayloadID = uUIDInfo.PayloadID
-			if _, err := database.DB.NamedExec(`INSERT INTO staginginfo 
-		(session_id, enc_key, dec_key, crypto_type, payload_id, staging_uuid)
-		VALUES (:session_id, :enc_key, :dec_key, :crypto_type, :payload_id, :staging_uuid)`, stagingDatabaseMessage); err != nil {
-				logging.LogError(err, "Failed to save staging information into database", "staginginfo", stagingDatabaseMessage)
-				return nil, errors.New(fmt.Sprintf("Failed to save staging information: %s", err.Error()))
-			} else {
-				// generate the response map
-				response := map[string]interface{}{}
-				response["uuid"] = tempUUID.String()
-				response["session_id"] = agentMessage.SessionID
-				response["session_key"] = base64.StdEncoding.EncodeToString(encryptedNewKey)
-				reflectBackOtherKeys(&response, &agentMessage.Other)
-				return response, nil
-			}
-		}
+		return nil, fmt.Errorf("failed to decode agent message into struct: %v", err)
+	}
+
+	newKey, err := mythicCrypto.GenerateKeysForPayload("aes256_hmac")
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate new AES key for staging rsa: %v", err)
+	}
+
+	publicKeyToUse := agentMessage.PublicKey
+	if !strings.HasPrefix(agentMessage.PublicKey, "LS0t") {
+		publicKeyToUse = fmt.Sprintf("-----BEGIN PUBLIC KEY-----\n%s\n-----END PUBLIC KEY-----", agentMessage.PublicKey)
+	}
+
+	decodedBytes, err := base64.StdEncoding.DecodeString(publicKeyToUse)
+	if err != nil {
+		return nil, fmt.Errorf("failed to base64 provided public key: %v", err)
+	}
+
+	publicKeyToUse = string(decodedBytes)
+	encryptedNewKey, err := mythicCrypto.RsaEncryptBytes(*newKey.EncKey, []byte(publicKeyToUse))
+	if err != nil {
+		return nil, fmt.Errorf("failed to encrypt new encryption key with RSA: %v", err)
+	}
+
+	tempUUID, err := uuid.NewRandom()
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate a new random UUID for staging: %v", err)
 	}
+
+	stagingDatabaseMessage.CryptoType = "aes256_hmac"
+	stagingDatabaseMessage.EncKey = newKey.EncKey
+	stagingDatabaseMessage.DecKey = newKey.DecKey
+	stagingDatabaseMessage.SessionID = agentMessage.SessionID
+	stagingDatabaseMessage.StagingUuID = tempUUID.String()
+	stagingDatabaseMessage.PayloadID = uUIDInfo.PayloadID
+
+	if _, err := database.DB.NamedExec(insertQuery, stagingDatabaseMessage); err != nil {
+		return nil, fmt.Errorf("failed to save staging information into database %s: %v", "staginginfo", err)
+	}
+	// generate the response map
+	response := map[string]interface{}{}
+	response["uuid"] = tempUUID.String()
+	response["session_id"] = agentMessage.SessionID
+	response["session_key"] = base64.StdEncoding.EncodeToString(encryptedNewKey)
+
+	reflectBackOtherKeys(&response, &agentMessage.Other)
+	return response, nil
 }