Skip to content

Commit

Permalink
Provide a secure way for register credential api
Browse files Browse the repository at this point in the history
  • Loading branch information
cb-github-robot committed Aug 21, 2024
2 parents a01414b + 8334dd9 commit 11537be
Show file tree
Hide file tree
Showing 9 changed files with 470 additions and 46 deletions.
51 changes: 48 additions & 3 deletions init/init.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,12 @@
from tqdm import tqdm
from colorama import Fore, init
from getpass import getpass
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA256
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad

parser = argparse.ArgumentParser(description="Automatically proceed without confirmation.")
parser.add_argument('-y', '--yes', action='store_true', help='Automatically answer yes to prompts and proceed.')
Expand Down Expand Up @@ -119,22 +125,61 @@ def get_decryption_key():

print(Fore.YELLOW + f"\nRegistering all valid credentials for all cloud regions...")

# Function to register credentials
# Function to encrypt credentials using AES and RSA public key
def encrypt_credential_value_with_publickey(public_key_pem, credentials):
public_key = RSA.import_key(public_key_pem)
rsa_cipher = PKCS1_OAEP.new(public_key, hashAlgo=SHA256)
aes_key = get_random_bytes(32) # AES-256 key

encrypted_credentials = {}
for k, v in credentials.items():
# Encrypt using AES
aes_cipher = AES.new(aes_key, AES.MODE_CBC)
ciphertext = aes_cipher.encrypt(pad(v.encode(), AES.block_size))
encrypted_credentials[k] = base64.b64encode(aes_cipher.iv + ciphertext).decode()

# Encrypt AES key with RSA and encode in Base64
encrypted_aes_key = base64.b64encode(rsa_cipher.encrypt(aes_key)).decode()

# Clear AES key from memory
del aes_key

return encrypted_credentials, encrypted_aes_key

# Function to register credentials using encrypted values and encrypted AES key
def register_credential(provider, credentials):
try:
if all(credentials.values()):
# Step 1: Get the public key for encryption
public_key_response = requests.get(f"http://{TUMBLEBUG_SERVER}/tumblebug/credential/publicKey", headers=HEADERS)
if public_key_response.status_code != 200:
return provider, "Failed to retrieve public key, Skip", Fore.RED

public_key_data = public_key_response.json()
public_key = public_key_data['publicKey']
public_key_token_id = public_key_data['publicKeyTokenId']

# Step 2: Encrypt the credentials using AES and RSA public key
encrypted_credentials, encrypted_aes_key = encrypt_credential_value_with_publickey(public_key, credentials)

# Step 3: Prepare the payload with the encrypted credentials and AES key
credential_payload = {
"credentialHolder": "admin",
"keyValueInfoList": [{"key": k, "value": v} for k, v in credentials.items()],
"providerName": provider
"credentialKeyValueList": [{"key": k, "value": v} for k, v in encrypted_credentials.items()],
"providerName": provider,
"publicKeyTokenId": public_key_token_id,
"encryptedAesKey": encrypted_aes_key
}

# Step 4: Register the encrypted credentials
response = requests.post(f"http://{TUMBLEBUG_SERVER}/tumblebug/credential", json=credential_payload, headers=HEADERS)
return provider, response.json(), Fore.GREEN
else:
return provider, "Incomplete credential data, Skip", Fore.RED
except Exception as e:
return provider, f"Error registering credentials: {str(e)}", Fore.RED


# Register credentials to TumblebugServer using ThreadPoolExecutor
with ThreadPoolExecutor(max_workers=5) as executor:
future_to_provider = {executor.submit(register_credential, provider, credentials): provider for provider, credentials in cred_data.items()}
Expand Down
5 changes: 4 additions & 1 deletion init/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -8,4 +8,7 @@ PyYAML==6.0.1
tqdm==4.66.3

# Used for coloring and styling console output
colorama==0.4.6
colorama==0.4.6

# Used for cryptographic functions like RSA encryption
pycryptodome==3.18.0
88 changes: 79 additions & 9 deletions src/api/rest/docs/docs.go
Original file line number Diff line number Diff line change
Expand Up @@ -449,7 +449,7 @@ const docTemplate = `{
"application/json"
],
"tags": [
"[Admin] Multi-Cloud environment configuration"
"[Admin] Credential Management"
],
"summary": "List all registered ConnConfig",
"operationId": "GetConnConfigList",
Expand Down Expand Up @@ -516,7 +516,7 @@ const docTemplate = `{
"application/json"
],
"tags": [
"[Admin] Multi-Cloud environment configuration"
"[Admin] Credential Management"
],
"summary": "Get registered ConnConfig info",
"operationId": "GetConnConfig",
Expand Down Expand Up @@ -553,17 +553,17 @@ const docTemplate = `{
},
"/credential": {
"post": {
"description": "Post register Credential info",
"description": "This API registers credential information using hybrid encryption. 1. First, compress and encrypt sensitive data using a client generated AES with a 256-bit key. 2. Then, encrypt the AES key using an RSA public key obtained from ` + "`" + `GET /credential/publicKey` + "`" + `. 3. RSA encryption uses a 4096-bit key with OAEP padding and SHA-256 as the hash function. Ensure that all values are base64 encoded before sending them in the request. The public key token ID must be included in the request to allow the server to decrypt the data.",
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"[Admin] Multi-Cloud environment configuration"
"[Admin] Credential Management"
],
"summary": "Post register Credential info",
"summary": "Register Credential Information",
"operationId": "RegisterCredential",
"parameters": [
{
Expand Down Expand Up @@ -598,6 +598,36 @@ const docTemplate = `{
}
}
},
"/credential/publicKey": {
"get": {
"description": "Generates an RSA key pair using a 4096-bit key size with the RSA algorithm. The public key is generated using the RSA algorithm with OAEP padding and SHA-256 as the hash function. This key is used to encrypt an AES key that will be used for hybrid encryption of credentials.",
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"[Admin] Credential Management"
],
"summary": "Get RSA Public Key for Credential Encryption",
"operationId": "GetPublicKeyForCredentialEncryption",
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/common.PublicKeyResponse"
}
},
"500": {
"description": "Internal Server Error",
"schema": {
"$ref": "#/definitions/common.SimpleMsg"
}
}
}
}
},
"/forward/{path}": {
"post": {
"description": "Forward any (GET) request to CB-Spider",
Expand Down Expand Up @@ -9048,19 +9078,35 @@ const docTemplate = `{
}
},
"common.CredentialReq": {
"description": "CredentialReq contains the necessary information to register a credential. This includes the AES key encrypted with the RSA public key, which is then used to decrypt the AES key on the server side.",
"type": "object",
"properties": {
"credentialHolder": {
"type": "string"
"description": "CredentialHolder is the entity or user that holds the credential.",
"type": "string",
"example": "admin"
},
"keyValueInfoList": {
"credentialKeyValueList": {
"description": "CredentialKeyValueList contains key-(encrypted)value pairs that include the sensitive credential data.",
"type": "array",
"items": {
"$ref": "#/definitions/common.KeyValue"
"$ref": "#/definitions/common.KeyWithEncryptedValue"
}
},
"encryptedClientAesKeyByPublicKey": {
"description": "EncryptedClientAesKeyByPublicKey is the client temporary AES key encrypted with the RSA public key.",
"type": "string",
"example": "ZzXL27hbAUDT0ohglf2Gwr60sAqdPw3+CnCsn0RJXeiZxXnHfW03mFx5RaSfbwtPYCq1h6wwv7XsiWzfFmr02..."
},
"providerName": {
"type": "string"
"description": "ProviderName specifies the cloud provider associated with the credential (e.g., AWS, GCP).",
"type": "string",
"example": "aws"
},
"publicKeyTokenId": {
"description": "PublicKeyTokenId is the unique token ID used to retrieve the corresponding private key for decryption.",
"type": "string",
"example": "cr31av30uphc738d7h0g"
}
}
},
Expand Down Expand Up @@ -9248,6 +9294,19 @@ const docTemplate = `{
}
}
},
"common.KeyWithEncryptedValue": {
"type": "object",
"properties": {
"key": {
"description": "Key for the value",
"type": "string"
},
"value": {
"description": "Should be encrypted by the public key issued by GET /credential/publicKey",
"type": "string"
}
}
},
"common.Location": {
"type": "object",
"properties": {
Expand Down Expand Up @@ -9292,6 +9351,17 @@ const docTemplate = `{
}
}
},
"common.PublicKeyResponse": {
"type": "object",
"properties": {
"publicKey": {
"type": "string"
},
"publicKeyTokenId": {
"type": "string"
}
}
},
"common.RegionDetail": {
"type": "object",
"properties": {
Expand Down
88 changes: 79 additions & 9 deletions src/api/rest/docs/swagger.json
Original file line number Diff line number Diff line change
Expand Up @@ -442,7 +442,7 @@
"application/json"
],
"tags": [
"[Admin] Multi-Cloud environment configuration"
"[Admin] Credential Management"
],
"summary": "List all registered ConnConfig",
"operationId": "GetConnConfigList",
Expand Down Expand Up @@ -509,7 +509,7 @@
"application/json"
],
"tags": [
"[Admin] Multi-Cloud environment configuration"
"[Admin] Credential Management"
],
"summary": "Get registered ConnConfig info",
"operationId": "GetConnConfig",
Expand Down Expand Up @@ -546,17 +546,17 @@
},
"/credential": {
"post": {
"description": "Post register Credential info",
"description": "This API registers credential information using hybrid encryption. 1. First, compress and encrypt sensitive data using a client generated AES with a 256-bit key. 2. Then, encrypt the AES key using an RSA public key obtained from `GET /credential/publicKey`. 3. RSA encryption uses a 4096-bit key with OAEP padding and SHA-256 as the hash function. Ensure that all values are base64 encoded before sending them in the request. The public key token ID must be included in the request to allow the server to decrypt the data.",
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"[Admin] Multi-Cloud environment configuration"
"[Admin] Credential Management"
],
"summary": "Post register Credential info",
"summary": "Register Credential Information",
"operationId": "RegisterCredential",
"parameters": [
{
Expand Down Expand Up @@ -591,6 +591,36 @@
}
}
},
"/credential/publicKey": {
"get": {
"description": "Generates an RSA key pair using a 4096-bit key size with the RSA algorithm. The public key is generated using the RSA algorithm with OAEP padding and SHA-256 as the hash function. This key is used to encrypt an AES key that will be used for hybrid encryption of credentials.",
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"[Admin] Credential Management"
],
"summary": "Get RSA Public Key for Credential Encryption",
"operationId": "GetPublicKeyForCredentialEncryption",
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/common.PublicKeyResponse"
}
},
"500": {
"description": "Internal Server Error",
"schema": {
"$ref": "#/definitions/common.SimpleMsg"
}
}
}
}
},
"/forward/{path}": {
"post": {
"description": "Forward any (GET) request to CB-Spider",
Expand Down Expand Up @@ -9041,19 +9071,35 @@
}
},
"common.CredentialReq": {
"description": "CredentialReq contains the necessary information to register a credential. This includes the AES key encrypted with the RSA public key, which is then used to decrypt the AES key on the server side.",
"type": "object",
"properties": {
"credentialHolder": {
"type": "string"
"description": "CredentialHolder is the entity or user that holds the credential.",
"type": "string",
"example": "admin"
},
"keyValueInfoList": {
"credentialKeyValueList": {
"description": "CredentialKeyValueList contains key-(encrypted)value pairs that include the sensitive credential data.",
"type": "array",
"items": {
"$ref": "#/definitions/common.KeyValue"
"$ref": "#/definitions/common.KeyWithEncryptedValue"
}
},
"encryptedClientAesKeyByPublicKey": {
"description": "EncryptedClientAesKeyByPublicKey is the client temporary AES key encrypted with the RSA public key.",
"type": "string",
"example": "ZzXL27hbAUDT0ohglf2Gwr60sAqdPw3+CnCsn0RJXeiZxXnHfW03mFx5RaSfbwtPYCq1h6wwv7XsiWzfFmr02..."
},
"providerName": {
"type": "string"
"description": "ProviderName specifies the cloud provider associated with the credential (e.g., AWS, GCP).",
"type": "string",
"example": "aws"
},
"publicKeyTokenId": {
"description": "PublicKeyTokenId is the unique token ID used to retrieve the corresponding private key for decryption.",
"type": "string",
"example": "cr31av30uphc738d7h0g"
}
}
},
Expand Down Expand Up @@ -9241,6 +9287,19 @@
}
}
},
"common.KeyWithEncryptedValue": {
"type": "object",
"properties": {
"key": {
"description": "Key for the value",
"type": "string"
},
"value": {
"description": "Should be encrypted by the public key issued by GET /credential/publicKey",
"type": "string"
}
}
},
"common.Location": {
"type": "object",
"properties": {
Expand Down Expand Up @@ -9285,6 +9344,17 @@
}
}
},
"common.PublicKeyResponse": {
"type": "object",
"properties": {
"publicKey": {
"type": "string"
},
"publicKeyTokenId": {
"type": "string"
}
}
},
"common.RegionDetail": {
"type": "object",
"properties": {
Expand Down
Loading

0 comments on commit 11537be

Please sign in to comment.