AirlockTools/services/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
import sys
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

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

logger = logging.getLogger(__name__)


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)

       
    logger.debug(f"API key for service '{service}' and user '{username}' stored successfully.")

    print("\n✅ API key stored securely.")
    print("The program will now exit. Press Enter to continue...")

    try:
        _ = input()
    except Exception:
        pass

    _ = None
    sys.exit(0)


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.")
        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