added support for encrypted credentials

pylo/API/APIConnector.py

@@ -2,6 +2,7 @@
 import time
 import getpass
 
+from .CredentialsManager import is_api_key_encrypted, decrypt_api_key
 from .JsonPayloadTypes import LabelGroupObjectJsonStructure, LabelObjectCreationJsonStructure, \
     LabelObjectJsonStructure, LabelObjectUpdateJsonStructure, PCEObjectsJsonStructure, \
     LabelGroupObjectUpdateJsonStructure, IPListObjectCreationJsonStructure, IPListObjectJsonStructure, \
@@ -65,14 +66,25 @@ def __init__(self, fqdn: str, port, apiuser: str, apikey: str, skip_ssl_cert_che
         if type(port) is int:
             port = str(port)
         self.port: int = port
-        self.api_key: str = apikey
+        self._api_key: str = apikey
+        self._decrypted_api_key: str = None
         self.api_user: str = apiuser
         self.orgID: int = org_id
         self.skipSSLCertCheck: bool = skip_ssl_cert_check
         self.version: Optional['pylo.SoftwareVersion'] = None
         self.version_string: str = "Not Defined"
         self._cached_session = requests.session()
 
+    @property
+    def api_key(self):
+        if self._decrypted_api_key is not None:
+            return self._decrypted_api_key
+        if is_api_key_encrypted(self._api_key):
+            self._decrypted_api_key = decrypt_api_key(self._api_key)
+            return self._decrypted_api_key
+        return self._api_key
+
+
     @staticmethod
     def get_all_object_types_names_except(exception_list: List[ObjectTypes]):
 

pylo/API/CredentialsManager.py

@@ -1,6 +1,9 @@
+import base64
+from hashlib import sha256
 from typing import Dict, TypedDict, Union, List, Optional
 import json
 import os
+from cryptography.fernet import Fernet
 from ..Exception import PyloEx
 from .. import log
 
@@ -148,7 +151,14 @@ def get_all_credentials() -> List[CredentialProfile]:
     return credentials
 
 
-def create_credential_in_file(file_full_path: str, data: CredentialFileEntry, overwrite_existing_profile = False) -> None:
+def create_credential_in_file(file_full_path: str, data: CredentialFileEntry, overwrite_existing_profile = False) -> str:
+    """
+    Create a credential in a file and return the full path to the file
+    :param file_full_path:
+    :param data:
+    :param overwrite_existing_profile:
+    :return:
+    """
     # if file already exists, load it and append the new credential to it
     if os.path.isdir(file_full_path):
         file_full_path = os.path.join(file_full_path, "credentials.json")
@@ -181,5 +191,95 @@ def create_credential_in_file(file_full_path: str, data: CredentialFileEntry, ov
     with open(file_full_path, 'w') as f:
         json.dump(credentials, f, indent=4)
 
-def create_credential_in_default_file(data: CredentialFileEntry) -> None:
-    create_credential_in_file(os.path.expanduser("~/.pylo/credentials.json"), data)
+    return file_full_path
+
+def create_credential_in_default_file(data: CredentialFileEntry) -> str:
+    """
+    Create a credential in the default credential file and return the full path to the file
+    :param data:
+    :return:
+    """
+    file_path = os.path.expanduser("~/.pylo/credentials.json")
+    create_credential_in_file(os.path.expanduser(file_path), data)
+
+
+def encrypt_api_key_with_paramiko_key(ssh_key: paramiko.AgentKey, api_key: str) -> str:
+
+
+    def encrypt(raw: str, key: bytes) -> bytes:
+        """
+
+        :param raw:
+        :param key:
+        :return: base64 encoded encrypted string
+        """
+        f = Fernet(base64.urlsafe_b64encode(key))
+        token = f.encrypt(bytes(raw, 'utf-8'))
+        return token
+
+
+    # generate a random 128bit key
+    session_key_to_sign = os.urandom(16)
+
+    signed_message = ssh_key.sign_ssh_data(session_key_to_sign)
+
+    # use SHA256 to hash the signed message and use it as final AES 256 key
+    encryption_key = sha256(signed_message).digest()
+    #print("Encryption key: {}".format(encryption_key.hex()))
+    encrypted_text = encrypt(api_key, encryption_key)
+
+    api_key = "$encrypted$:ssh-Fernet:{}:{}:{}".format(base64.urlsafe_b64encode(ssh_key.get_fingerprint()).decode('utf-8'),
+                                                       base64.urlsafe_b64encode(session_key_to_sign).decode('utf-8'),
+                                                       encrypted_text.decode('utf-8'))
+
+    return api_key
+
+
+def decrypt_api_key_with_paramiko_key(encrypted_api_key_payload: str) -> str:
+    def decrypt(token_b64_encoded: str, key: bytes):
+        f = Fernet(base64.urlsafe_b64encode(key))
+        return f.decrypt(token_b64_encoded).decode('utf-8')
+
+    # split the api_key into its components
+    api_key_parts = encrypted_api_key_payload.split(":")
+    if len(api_key_parts) != 5:
+        raise PyloEx("Invalid encrypted API key format")
+
+    # get the fingerprint and the session key
+    fingerprint = base64.urlsafe_b64decode(api_key_parts[2])
+    session_key = base64.urlsafe_b64decode(api_key_parts[3])
+    encrypted_api_key = api_key_parts[4]
+
+    # find the key in the agent
+    keys = paramiko.Agent().get_keys()
+    found_key = None
+    for key in keys:
+        if key.get_fingerprint() == fingerprint:
+            found_key = key
+            break
+
+    if found_key is None:
+        raise PyloEx("No key found in the agent with fingerprint {}".format(fingerprint.hex()))
+
+    # sign the session key
+    signed_session_key = found_key.sign_ssh_data(session_key)
+    encryption_key = sha256(signed_session_key).digest()
+    #print("Encryption key: {}".format(encryption_key.hex()))
+    #print("Encrypted from KEY fingerprint: {}".format(fingerprint.hex()))
+
+    return decrypt(token_b64_encoded=encrypted_api_key,
+                   key=encryption_key
+                   )
+
+def decrypt_api_key(encrypted_api_key_payload: str) -> str:
+    # detect the encryption method
+    if not encrypted_api_key_payload.startswith("$encrypted$:"):
+        raise PyloEx("Invalid encrypted API key format")
+    if encrypted_api_key_payload.startswith("$encrypted$:ssh-Fernet:"):
+        return decrypt_api_key_with_paramiko_key(encrypted_api_key_payload)
+
+    raise PyloEx("Unsupported encryption method: {}".format(encrypted_api_key_payload.split(":")[1]))
+
+
+def is_api_key_encrypted(encrypted_api_key_payload: str) -> bool:
+    return encrypted_api_key_payload.startswith("$encrypted$:")

pylo/cli/commands/credential_manager.py

@@ -1,11 +1,13 @@
-import logging
+from prettytable import PrettyTable
 import argparse
 import os
 
+import paramiko
+
 import pylo
 import click
 from pylo.API.CredentialsManager import get_all_credentials, create_credential_in_file, CredentialFileEntry, \
-    create_credential_in_default_file
+    create_credential_in_default_file, encrypt_api_key_with_paramiko_key, decrypt_api_key_with_paramiko_key
 
 from pylo import log
 from . import Command
@@ -81,25 +83,80 @@ def __main(args, **kwargs):
             "api_key": api_key
         }
 
+        encrypt_api_key = click.prompt('> Encrypt API? Y/N', type=bool)
+        if encrypt_api_key:
+            print("Available keys (ECDSA NISTPXXX keys and a few others are not supported and will be filtered out):")
+            ssh_keys = paramiko.Agent().get_keys()
+            # filter out ECDSA NISTPXXX and [email protected]
+            ssh_keys = [key for key in ssh_keys if not (key.get_name().startswith("ecdsa-sha2-nistp") or
+                                                        key.get_name().startswith("[email protected]"))
+                                                        ]
+
+            # display a table of keys
+            print_keys(keys=ssh_keys, display_index=True)
+            print()
+
+            index_of_selected_key = click.prompt('> Select key by ID#', type=click.IntRange(0, len(ssh_keys)-1))
+            selected_ssh_key = ssh_keys[index_of_selected_key]
+            print("Selected key: {} | {} | {}".format(selected_ssh_key.get_name(),
+                                                      selected_ssh_key.get_fingerprint().hex(),
+                                                      selected_ssh_key.comment))
+            print(" * encrypting API key with selected key...", flush=True, end="")
+            encrypted_api_key = encrypt_api_key_with_paramiko_key(ssh_key=selected_ssh_key, api_key=api_key)
+            print("OK!")
+            print(" * trying to decrypt the encrypted API key...", flush=True, end="")
+            decrypted_api_key = decrypt_api_key_with_paramiko_key(encrypted_api_key_payload=encrypted_api_key)
+            if decrypted_api_key != api_key:
+                raise pylo.PyloEx("Decrypted API key does not match original API key")
+            print("OK!")
+            credentials_data["api_key"] = encrypted_api_key
+
+
         cwd = os.getcwd()
-        create_in_current_workdir = click.prompt('> Create in current workdir? Y/N ({})'.format(cwd), type=bool)
+        create_in_current_workdir = click.prompt('> Create in current workdir ({})? If not then user homedir will be used.   Y/N '.format(cwd), type=bool)
 
 
         print("* Creating credential...", flush=True, end="")
         if create_in_current_workdir:
-            create_credential_in_file(file_full_path=cwd, data=credentials_data)
+            file_path = create_credential_in_file(file_full_path=cwd, data=credentials_data)
         else:
-            create_credential_in_default_file(data=credentials_data)
+            file_path = create_credential_in_default_file(data=credentials_data)
+
+        print("OK! ({})".format(file_path))
 
-        print("OK!")
 
+command_object = Command(command_name, __main, fill_parser, credentials_manager_mode=True)
 
+def print_keys(keys: list[paramiko.AgentKey], display_index = True) -> None:
 
+    args_for_print = []
 
+    column_properties = [  # (name, width)
+        ("ID#", 4),
+        ("Name", 20),
+        ("Fingerprint", 40),
+        ("Comment", 48)
+    ]
 
+    if not display_index:
+        # remove tuple with name "ID#"
+        column_properties = [item for item in column_properties if item[0] != "ID#"]
 
 
+    table = PrettyTable()
+    table.field_names = [item[0] for item in column_properties]
 
 
-command_object = Command(command_name, __main, fill_parser, credentials_manager_mode=True)
+    for i, key in enumerate(keys):
+        display_values = []
+        if display_index:
+            display_values.append(i)
+        display_values.append(key.get_name())
+        display_values.append(key.get_fingerprint().hex())
+        display_values.append(key.comment)
+
+        table.add_row(display_values)
+
+    print(table)
+
 

requirements.txt

@@ -1,6 +1,8 @@
 click==8.1.7
 colorama~=0.4.4
-openpyxl~=3.0.10
+cryptography~=42.0.5
+openpyxl~=3.1.2
 paramiko~=3.4.0
+prettytable~=3.10.0
 requests~=2.31.0
 xlsxwriter~=1.3.7

setup.py

@@ -23,8 +23,10 @@
     install_requires=[
         'click~=8.1.7',
         'colorama~=0.4.4',
+        'cryptography~=42.0.5',
         'openpyxl~=3.0.10',
         'paramiko~=3.4.0',
+        'prettytable~=3.10.0'
         'requests~=2.31.0',
         'xlsxwriter~=1.3.7',
     ],