AirlockTools/Development/WIP/Sync Code/security.py

#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

import base64
import logging
import os
import platform
import re
from getpass import getpass

import keyring
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from utils.utils import colorText
from sys import exit

# Constants
KDF_ITERATIONS = 200_000
SALT_SIZE = 16 # 128-bit Salt
NONCE_SIZE = 12 # AES-GCM 
KEY_SIZE = 32  # AES-256


def _derive_key(password: bytes, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=salt,
        iterations=KDF_ITERATIONS,
    )
    return kdf.derive(password)


def configure_keyring_backend():
    system = platform.system()
    if system == "Windows":
        import keyring.backends.Windows
        keyring.set_keyring(keyring.backends.Windows.WinVaultKeyring())
    elif system == "Linux":
        import keyring.backends.kwallet
        keyring.set_keyring(keyring.backends.kwallet.DBusKeyring())
    else:
        raise EnvironmentError(f"Unsupported OS: {system}")


def store_api_key(service: str, username: str, api_key: str, password: str):
    configure_keyring_backend()
    salt = os.urandom(SALT_SIZE)
    key = _derive_key(password.encode(), salt)
    aesgcm = AESGCM(key)
    nonce = os.urandom(NONCE_SIZE)
    ct = aesgcm.encrypt(nonce, api_key.encode(), associated_data=None)
    blob = salt + nonce + ct
    b64 = base64.b64encode(blob).decode()
    keyring.set_password(service, username, b64)


def retrieve_api_key(service: str, username: str, password: str) -> str:
    configure_keyring_backend()
    b64 = keyring.get_password(service, username)
    if b64 is None:
        raise ValueError("No stored secret for this service/username.")
    blob = base64.b64decode(b64)
    salt = blob[:SALT_SIZE]
    nonce = blob[SALT_SIZE:SALT_SIZE + NONCE_SIZE]
    ct = blob[SALT_SIZE + NONCE_SIZE:]
    key = _derive_key(password.encode(), salt)
    aesgcm = AESGCM(key)
    pt = aesgcm.decrypt(nonce, ct, associated_data=None)
    return pt.decode()


def api_key_exists(service: str, username: str) -> bool:
    configure_keyring_backend()
    return keyring.get_password(service, username) is not None


def check_password_complexity(password: str) -> bool:
    if len(password) < 12:
        return False
    if not re.search(r"[A-Z]", password):
        return False
    if not re.search(r"[a-z]", password):
        return False
    if not re.search(r"[0-9]", password):
        return False
    if not re.search(r"[^A-Za-z0-9]", password):
        return False
    return True


def getAPI(USERNAME, SERVICE_NAME):
    logging.debug(
        f"Checking for stored API key for user '{USERNAME}' in service '{SERVICE_NAME}'..."
    )

    if api_key_exists(SERVICE_NAME, USERNAME):
        for attempt in range(1, 4):
            password = getpass(f"Attempt {attempt}/3 - Enter password to unlock your API key: ")
            try:
                apikey = retrieve_api_key(SERVICE_NAME, USERNAME, password)
                logging.debug("API key successfully retrieved.")
                return apikey
            except Exception as e:
                logging.warning(f"Attempt {attempt} failed: {str(e)}")
        logging.error("Failed to retrieve API key after 3 incorrect attempts.")      
        print(colorText("❌ Authentication failed. Exiting.","red"))
        exit(1)  # Exit cleanly without traceback
        raise ValueError("Failed to retrieve API key after 3 incorrect attempts.")
    else:
        logging.warning(f"No API key found for user '{USERNAME}' in service '{SERVICE_NAME}'.")
        api_key = getpass(f"🗝️ No API key found. Please enter your API key for '{SERVICE_NAME}': ").strip()
        print("Please exit and relaunch program after saving your credential to avoid errors")

        while True:
            password = getpass("🔓 Create a password to encrypt your API key: ")
            confirm_password = getpass("🔒 Confirm your password: ")

            if password != confirm_password:
                logging.warning("❌ Passwords do not match. Try again.")
                continue

            if check_password_complexity(password):
                try:
                    store_api_key(SERVICE_NAME, USERNAME, api_key, password)
                    logging.info("API key stored securely.")
                    break
                except Exception as e:
                    logging.error(f"Failed to store API key: {e}")
                    break
            else:
                logging.warning("❌ Password does not meet complexity requirements. Try again.")


class APIKeyManager:
    _api_key = None

    @classmethod
    def load(cls, service: str, username: str, password: str):
        cls._api_key = retrieve_api_key(service, username, password)

    @classmethod
    def get(cls) -> str:
        if cls._api_key is None:
            raise ValueError("API key not loaded. Call APIKeyManager.load() first.")
        return cls._api_key