139 lines
4.8 KiB
Python
139 lines
4.8 KiB
Python
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
|
|
|
|
# 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}")
|
|
|
|
|
|
async 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)
|
|
|
|
|
|
async 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()
|
|
|
|
|
|
async 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:
|
|
return (
|
|
len(password) >= 12
|
|
and bool(re.search(r"[A-Z]", password))
|
|
and bool(re.search(r"[a-z]", password))
|
|
and bool(re.search(r"[0-9]", password))
|
|
and bool(re.search(r"[^A-Za-z0-9]", password))
|
|
)
|
|
|
|
|
|
async def getAPI(USERNAME, SERVICE_NAME):
|
|
logging.debug(
|
|
f"Checking for stored API key for user '{USERNAME}' in service '{SERVICE_NAME}'..."
|
|
)
|
|
|
|
if await 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 = await 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)
|
|
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:
|
|
await 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
|
|
async def load(cls, service: str, username: str, password: str):
|
|
cls._api_key = await 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 |