#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
AEA Smart Gateway - Distributed Database Encryption Module (Relay Mode)
Developer: AEA Team (Version: 1.24)
=============================================================================
Provides split-key encryption for the database file using AES-256-CBC.
Key structure (60 characters):
- 40 characters (Head): Stored locally in /etc/aea_gateway_db.key [Root Only]
- 20 characters (Tail): Generated dynamically and appended to the .enc file
"""

import os
import sys
import secrets
import subprocess

# --- DYNAMIC PATH MACROS (REPLACED DURING DEPLOYMENT) ---
KEY_FILE_PATH = "{{CRYPTO_KEY_PATH}}"
DB_FILE_PATH = "{{CRYPTO_DB_PATH}}"

# --- STATUS AND ERROR STRINGS FOR SYSTEM LOGGER AND ANDROID UI ---
TXT_ERR_NO_KEY_FILE = "[ERR_CRYPTO_01] System key file (40 chars) not found!"
TXT_ERR_NO_DB_FILE  = "[ERR_CRYPTO_02] Encrypted database file not found!"
TXT_ERR_DB_CORRUPT  = "[ERR_CRYPTO_03] Database corrupted (size less than 20 bytes tail)!"
TXT_ERR_DECRYPT     = "[ERR_CRYPTO_04] Fatal decryption error! Invalid key or corrupted data."
TXT_ERR_OPENSSL     = "[ERR_CRYPTO_05] Critical system OpenSSL call error."
TXT_SUCCESS_INIT    = "[SUCCESS_CRYPTO] Gateway crypto environment initialized successfully."


def get_head_key():
    """Reads the first 40 characters of the key from the secure root directory."""
    if not os.path.exists(KEY_FILE_PATH):
        print(TXT_ERR_NO_KEY_FILE, file=sys.stderr)
        sys.exit(1)
    with open(KEY_FILE_PATH, "r", encoding="utf-8") as f:
        return f.read().strip()


def init_environment():
    """Initial environment setup during gateway deployment. Creates keys and empty DB."""
    os.makedirs(os.path.dirname(DB_FILE_PATH), exist_ok=True)
    os.makedirs(os.path.dirname(KEY_FILE_PATH), exist_ok=True)
    
    # 1. Generate permanent head key (40 hex chars = 20 bytes)
    if not os.path.exists(KEY_FILE_PATH):
        head_key = secrets.token_hex(20)
        with open(KEY_FILE_PATH, "w", encoding="utf-8") as f:
            f.write(head_key)
        os.chmod(KEY_FILE_PATH, 0o600)
    
    # 2. Create clean database file (empty JSON array '[]')
    write_database("[]")
    print(TXT_SUCCESS_INIT)


def read_database():
    """Decrypts the database, strips the dynamic tail, and outputs clean JSON string."""
    if not os.path.exists(DB_FILE_PATH):
        print(TXT_ERR_NO_DB_FILE, file=sys.stderr)
        sys.exit(2)
        
    with open(DB_FILE_PATH, "rb") as f:
        file_content = f.read()
        
    if len(file_content) < 20:
        print(TXT_ERR_DB_CORRUPT, file=sys.stderr)
        sys.exit(3)
        
    # Split file: encrypted payload and last 20 bytes of dynamic tail key
    encrypted_data = file_content[:-20]
    tail_key = file_content[-20:].decode('utf-8', errors='ignore')
    
    # Assemble full 60-character password
    head_key = get_head_key()
    full_60char_password = head_key + tail_key
    
    cmd = [
        "openssl", "enc", "-aes-256-cbc", "-d", "-pbkdf2", 
        "-iter", "100000", "-pass", f"pass:{full_60char_password}"
    ]
    
    process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate(input=encrypted_data)
    
    if process.returncode != 0:
        print(TXT_ERR_DECRYPT, file=sys.stderr)
        sys.exit(4)
        
    return stdout.decode('utf-8')


def write_database(raw_json_string):
    """Encrypts raw JSON string, generates a new tail key, and saves the payload."""
    os.makedirs(os.path.dirname(DB_FILE_PATH), exist_ok=True)
    
    # Generate new random tail key (20 hex chars = 10 bytes)
    tail_key = secrets.token_hex(10)
    head_key = get_head_key()
    full_60char_password = head_key + tail_key
    
    cmd = [
        "openssl", "enc", "-aes-256-cbc", "-salt", "-pbkdf2", 
        "-iter", "100000", "-pass", f"pass:{full_60char_password}"
    ]
    
    process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    encrypted_data, stderr = process.communicate(input=raw_json_string.encode('utf-8'))
    
    if process.returncode != 0:
        print(TXT_ERR_OPENSSL, file=sys.stderr)
        sys.exit(5)
        
    # Append 20 characters of unencrypted tail key to the end of binary encrypted data
    final_file_content = encrypted_data + tail_key.encode('utf-8')
    
    with open(DB_FILE_PATH, "wb") as f:
        f.write(final_file_content)
    
    os.chmod(DB_FILE_PATH, 0o600)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: aea_db_crypto.py [--init | --read | --write 'json_data']")
        sys.exit(6)
        
    action = sys.argv[1]
    
    if action == "--init":
        init_environment()
    elif action == "--read":
        print(read_database())
    elif action == "--write":
        if len(sys.argv) < 3:
            print("Error: Missing JSON payload for write operation!", file=sys.stderr)
            sys.exit(7)
        write_database(sys.argv[2])
    else:
        print(f"Unknown command: {action}", file=sys.stderr)
        sys.exit(8)