import base64
import datetime
import json
import os
import time
from os import mkdir
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor

from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.PublicKey import RSA
from Crypto.Util.Padding import unpad
from colorama import Fore
from colorama import init
from openpyxl import load_workbook
from tqdm import tqdm
from unidecode import unidecode

init()
current = datetime.datetime.now()
try:
    os.mkdir("output")
except FileExistsError:
    pass

def extract_system_info(lines):
    """Extract system information from decrypted file content."""
    system_info = {
        'script_version': '',
        'os_name': '',
        'kernel_version': '',
        'architecture': '',
        'hostname': '',
        'fqdn': '',
        'ip_address': '',
        'all_ips': '',
        'audit_time': '',
        'timezone': '',
        'uptime': '',
        'cpu_model': '',
        'cpu_cores': '',
        'total_memory': '',
        'used_memory': '',
        'free_memory': '',
        'swap_total': '',
        'primary_interface': '',
        'mac_address': '',
        'default_gateway': '',
        'dns_servers': '',
        'selinux_status': '',
        'firewall_status': '',
        'last_boot': '',
        'current_users': '',
        'load_average': '',
        # Compliance statistics
        'passed_count': 0,
        'failed_count': 0,
        'total_checks': 0,
        'compliance_percentage': 0.0,
        # Mandatory/Optional statistics
        'mandatory_passed': 0,
        'mandatory_failed': 0,
        'mandatory_total': 0,
        'mandatory_percentage': 0.0,
        'optional_passed': 0,
        'optional_failed': 0,
        'optional_total': 0,
        'optional_percentage': 0.0
    }
    
    passed_count = 0
    failed_count = 0
    
    for line in lines:
        line_stripped = line.strip()
        
        # Count PASSED and FAILED checks from JSON format lines
        if line_stripped.startswith('{') and line_stripped.endswith('}'):
            if '"PASSED"' in line_stripped:
                passed_count += 1
            elif '"FAILED"' in line_stripped:
                failed_count += 1
        
        # Parse system information
        if line_stripped.startswith('Script Version:'):
            system_info['script_version'] = line_stripped.replace('Script Version:', '').strip()
        elif line_stripped.startswith('Operating System:'):
            system_info['os_name'] = line_stripped.replace('Operating System:', '').strip()
        elif line_stripped.startswith('Kernel Version:'):
            system_info['kernel_version'] = line_stripped.replace('Kernel Version:', '').strip()
        elif line_stripped.startswith('Architecture:'):
            system_info['architecture'] = line_stripped.replace('Architecture:', '').strip()
        elif line_stripped.startswith('Hostname:'):
            system_info['hostname'] = line_stripped.replace('Hostname:', '').strip()
        elif line_stripped.startswith('FQDN:'):
            system_info['fqdn'] = line_stripped.replace('FQDN:', '').strip()
        elif line_stripped.startswith('IP Address:'):
            system_info['ip_address'] = line_stripped.replace('IP Address:', '').strip()
        elif line_stripped.startswith('All IP Addresses:'):
            system_info['all_ips'] = line_stripped.replace('All IP Addresses:', '').strip()
        elif line_stripped.startswith('Audit Time:'):
            system_info['audit_time'] = line_stripped.replace('Audit Time:', '').strip()
        elif line_stripped.startswith('Timezone:'):
            system_info['timezone'] = line_stripped.replace('Timezone:', '').strip()
        elif line_stripped.startswith('Uptime:'):
            system_info['uptime'] = line_stripped.replace('Uptime:', '').strip()
        elif line_stripped.startswith('CPU Model:'):
            system_info['cpu_model'] = line_stripped.replace('CPU Model:', '').strip()
        elif line_stripped.startswith('CPU Cores:'):
            system_info['cpu_cores'] = line_stripped.replace('CPU Cores:', '').strip()
        elif line_stripped.startswith('Total Memory:'):
            system_info['total_memory'] = line_stripped.replace('Total Memory:', '').strip()
        elif line_stripped.startswith('Used Memory:'):
            system_info['used_memory'] = line_stripped.replace('Used Memory:', '').strip()
        elif line_stripped.startswith('Free Memory:'):
            system_info['free_memory'] = line_stripped.replace('Free Memory:', '').strip()
        elif line_stripped.startswith('Swap Total:'):
            system_info['swap_total'] = line_stripped.replace('Swap Total:', '').strip()
        elif line_stripped.startswith('Primary Interface:'):
            system_info['primary_interface'] = line_stripped.replace('Primary Interface:', '').strip()
        elif line_stripped.startswith('MAC Address:'):
            system_info['mac_address'] = line_stripped.replace('MAC Address:', '').strip()
        elif line_stripped.startswith('Default Gateway:'):
            system_info['default_gateway'] = line_stripped.replace('Default Gateway:', '').strip()
        elif line_stripped.startswith('DNS Servers:'):
            system_info['dns_servers'] = line_stripped.replace('DNS Servers:', '').strip()
        elif line_stripped.startswith('SELinux Status:'):
            system_info['selinux_status'] = line_stripped.replace('SELinux Status:', '').strip()
        elif line_stripped.startswith('Firewall:'):
            system_info['firewall_status'] = line_stripped.replace('Firewall:', '').strip()
        elif line_stripped.startswith('Last Boot:'):
            system_info['last_boot'] = line_stripped.replace('Last Boot:', '').strip()
        elif line_stripped.startswith('Current Users:'):
            system_info['current_users'] = line_stripped.replace('Current Users:', '').strip()
        elif line_stripped.startswith('Load Average:'):
            system_info['load_average'] = line_stripped.replace('Load Average:', '').strip()
    
    # Calculate compliance statistics
    total_checks = passed_count + failed_count
    system_info['passed_count'] = passed_count
    system_info['failed_count'] = failed_count
    system_info['total_checks'] = total_checks
    
    if total_checks > 0:
        system_info['compliance_percentage'] = round((passed_count / total_checks) * 100, 1)
    else:
        system_info['compliance_percentage'] = 0.0
    
    return system_info

def detect_os_from_content(lines):
    """
    Detect OS type from decrypted file content.
    Reads the 'Operating System:' field and maps it to an OS key.
    
    Returns the OS key (centos, ubuntu, rhel, oracle, windows) or None if not detected.
    
    Example values:
        'CentOS Linux 7 (Core)'  -> 'centos'
        'Red Hat Enterprise ...' -> 'rhel'
        'Ubuntu 20.04.x LTS'    -> 'ubuntu'
        'Oracle Linux 8.x'      -> 'oracle'
        'Windows Server 2019'   -> 'windows'
    """
    os_name_raw = None
    for line in lines:
        stripped = line.strip()
        if stripped.startswith('Operating System:'):
            os_name_raw = stripped.replace('Operating System:', '').strip()
            break
    
    if not os_name_raw:
        return None
    
    os_lower = os_name_raw.lower()

    # Order matters: more specific patterns first
    if 'centos' in os_lower:
        detected = 'centos'
    elif 'red hat' in os_lower or 'redhat' in os_lower or 'rhel' in os_lower:
        detected = 'rhel'
    elif 'ubuntu' in os_lower:
        detected = 'ubuntu'
    elif 'oracle' in os_lower:
        # Check if version is < 8
        # Hỗ trợ cả 'Oracle Linux Server release 7.9' và 'Oracle Linux Server 7.9'
        import re
        version_match = re.search(r'(?:release\s+|server\s+|linux\s+)(\d+)', os_lower)
        if not version_match:
            # Fallback: tìm số version đầu tiên trong chuỗi (ví dụ 'oracle linux 7.9')
            version_match = re.search(r'(\d+)\.', os_lower)
        if version_match and int(version_match.group(1)) < 8:
            print(Fore.YELLOW + f"[CONTENT-DETECT] Oracle Linux < 8 detected (version {version_match.group(1)}). Routing to CentOS config." + Fore.RESET)
            detected = 'centos'
        else:
            detected = 'oracle'
    elif 'windows' in os_lower:
        detected = 'windows'
    else:
        detected = None

    if detected:
        print(Fore.CYAN + f"[CONTENT-DETECT] 'Operating System: {os_name_raw}' → OS key: '{detected}'" + Fore.RESET)
    else:
        print(Fore.YELLOW + f"[CONTENT-DETECT] Unknown OS from content: '{os_name_raw}'" + Fore.RESET)

    return detected


_cached_private_key = None

def get_private_key():
    """Lấy hoặc nạp RSA private key từ cache trong RAM."""
    global _cached_private_key
    if _cached_private_key is not None:
        return _cached_private_key

    script_dir = os.path.dirname(os.path.abspath(__file__))
    private_key_paths = [
        'private_key.pem',
        '../private_key.pem',
        os.path.join(script_dir, 'private_key.pem'),
        os.path.join(script_dir, 'keys', 'private_key.pem'),
        'keys/private_key.pem'
    ]

    for path in private_key_paths:
        if os.path.exists(path):
            with open(path, 'rb') as f:
                _cached_private_key = RSA.import_key(f.read())
                return _cached_private_key

    raise FileNotFoundError("private_key.pem not found! Please ensure it exists in the script directory.")


def _decrypt_rsa_block_worker(block: bytes) -> bytes:
    """Hàm worker cấp module phục vụ cho ProcessPoolExecutor / ThreadPoolExecutor giải mã RSA."""
    private_key = get_private_key()
    cipher = PKCS1_OAEP.new(private_key)
    return cipher.decrypt(block)


def is_hybrid_encrypted(raw_text: str) -> bool:
    """Kiểm tra xem nội dung file có thuộc chuẩn mã hóa lai MỚI (HYBRID AES-256 + RSA) hay không."""
    s = raw_text.strip()
    return s.startswith("HYBRID_V1:") or s.startswith("HYBRID_AES:") or s.startswith("HYBRID:")


def decrypt_hybrid(raw_text: str) -> bytes:
    """
    Giải mã định dạng Mã hóa Lai (Hybrid AES-256 + RSA) siêu tốc (~0.005s).
    Cấu trúc:
        HYBRID_V1:<b64_enc_key>:<b64_iv>:<b64_ciphertext>
    hoặc JSON:
        HYBRID_V1:{"enc_key": "...", "iv": "...", "data": "..."}
    """
    s = raw_text.strip()
    # Loại bỏ tiền tố prefix
    for prefix in ("HYBRID_V1:", "HYBRID_AES:", "HYBRID:"):
        if s.startswith(prefix):
            payload = s[len(prefix):].strip()
            break

    b64_enc_key = ""
    b64_iv = ""
    b64_ciphertext = ""

    if payload.startswith("{"):
        # Format JSON
        data_json = json.loads(payload)
        b64_enc_key = data_json.get("enc_key", "")
        b64_iv = data_json.get("iv", "")
        b64_ciphertext = data_json.get("data", "")
    else:
        # Format chuỗi phân cách bởi dấu hai chấm ':'
        parts = payload.split(":")
        if len(parts) >= 3:
            b64_enc_key = parts[0]
            b64_iv = parts[1]
            b64_ciphertext = parts[2]

    if not b64_enc_key or not b64_iv or not b64_ciphertext:
        raise ValueError("Cấu trúc file Mã hóa Lai (Hybrid Encryption) không hợp lệ!")

    enc_key = base64.b64decode(b64_enc_key)
    iv = base64.b64decode(b64_iv)
    ciphertext = base64.b64decode(b64_ciphertext)

    # 1. Giải mã khóa AES bằng RSA Private Key (chỉ 1 block RSA duy nhất)
    private_key = get_private_key()
    cipher_rsa = PKCS1_OAEP.new(private_key)
    aes_key = cipher_rsa.decrypt(enc_key)

    # 2. Giải mã dữ liệu bằng khóa AES-256-CBC (siêu nhanh)
    cipher_aes = AES.new(aes_key, AES.MODE_CBC, iv)
    decrypted_padded = cipher_aes.decrypt(ciphertext)

    # Thử gỡ padding PKCS7
    try:
        plaintext_bytes = unpad(decrypted_padded, AES.block_size)
    except ValueError:
        plaintext_bytes = decrypted_padded

    return plaintext_bytes


def decryption_read_header(filename: str, max_blocks: int = 30) -> list[str]:
    """
    Giải mã nhanh phần header của file .enc để đọc thông tin OS.
    Tự động hỗ trợ cả chuẩn MỚI (Hybrid AES+RSA) và chuẩn CŨ (Pure RSA).
    """
    with open(filename, "r", encoding="utf-8", errors="replace") as f:
        raw_text = f.read()

    if is_hybrid_encrypted(raw_text):
        # Format MỚI: Giải mã siêu tốc bằng AES-256 + RSA (< 0.005s)
        plaintext_bytes = decrypt_hybrid(raw_text)
        lines = plaintext_bytes.decode("utf-8", errors="replace").splitlines()
        return lines[:100]

    # Format CŨ: Giải mã N block RSA đầu tiên (~5KB)
    ciphertext = base64.b64decode(raw_text.encode('utf-8'))
    private_key = get_private_key()
    block_size = private_key.size_in_bytes()
    cipher = PKCS1_OAEP.new(private_key)

    total_blocks = (len(ciphertext) + block_size - 1) // block_size
    num_blocks = min(max_blocks, total_blocks)
    blocks = []

    for i in range(num_blocks):
        offset = i * block_size
        block = ciphertext[offset:offset + block_size]
        try:
            blocks.append(cipher.decrypt(block))
        except Exception:
            break

    header_text = b"".join(blocks).decode('utf-8', errors='replace')
    return header_text.splitlines()


def decryption_write(filename):
    """
    Giải mã toàn bộ file mã hóa.
    Tự động chuyển đổi giữa chuẩn MỚI (Hybrid AES-256 + RSA) và chuẩn CŨ (Pure RSA Multi-processing).
    """
    with open(filename, "r", encoding="utf-8", errors="replace") as f:
        raw_text = f.read()

    start_time = time.time()

    if is_hybrid_encrypted(raw_text):
        # ⚡ Chuẩn MỚI (Hybrid AES-256 + RSA): Tốc độ tức thì < 0.01 giây
        plaintext_bytes = decrypt_hybrid(raw_text)
        plaintext = plaintext_bytes.decode("utf-8", errors="replace")
        print(Fore.GREEN + f"  [HYBRID AES+RSA] Decryption done in {time.time() - start_time:.3f}s" + Fore.RESET)
    else:
        # 🐢 Chuẩn CŨ (Pure RSA): Giải mã đa luồng / đa tiến trình tối ưu CPU Multi-core
        ciphertext = base64.b64decode(raw_text.encode('utf-8'))
        private_key = get_private_key()
        block_size = private_key.size_in_bytes()
        total_blocks = (len(ciphertext) + block_size - 1) // block_size

        block_list = [ciphertext[i * block_size:(i + 1) * block_size] for i in range(total_blocks)]
        max_workers = min(16, (os.cpu_count() or 4) * 2)

        try:
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                decrypted_blocks = list(executor.map(_decrypt_rsa_block_worker, block_list))
        except Exception:
            cipher = PKCS1_OAEP.new(private_key)
            decrypted_blocks = [cipher.decrypt(b) for b in block_list]

        plaintext = b"".join(decrypted_blocks).decode('utf-8', errors='replace')
        print(Fore.GREEN + f"  [LEGACY RSA] Decryption done in {time.time() - start_time:.1f}s ({total_blocks} blocks, {max_workers} workers)" + Fore.RESET)

    output_path = os.path.basename(filename).replace(".enc", "")
    with open(output_path, 'w', encoding="utf-8") as f:
        f.write(plaintext)


def check_file_exist(filename):
    isFile = os.path.isfile(filename)
    if isFile:
        return True
    else:
        return False


def string_normalization(string):
    return unidecode(str(string))


def extract_failed_logs(lines):
    """Extract detailed log blocks for each FAILED criteria from decrypted log lines.
    
    Returns a dict mapping normalized criteria name -> log detail text.
    """
    failed_logs = {}
    i = 0
    separator = '######################################'
    dash_separator = '-------------------'
    stderr_begin = '########## STDERR BEGIN ##########'
    stderr_end = '########## STDERR END ##########'
    
    while i < len(lines):
        line_stripped = lines[i].strip()
        
        # Look for FAILED JSON lines
        if line_stripped.startswith('{') and line_stripped.endswith('}') and '"FAILED"' in line_stripped:
            try:
                parsed = json.loads(line_stripped)
                criteria_name = str(list(parsed.keys())[0])
                normalized_name = string_normalization(criteria_name)
                
                # Collect all lines after this FAILED line until the next JSON criteria line
                log_lines = []
                j = i + 1
                while j < len(lines):
                    next_line = lines[j].strip()
                    
                    # Stop when we hit the next JSON criteria line (PASSED or FAILED)
                    if next_line.startswith('{') and next_line.endswith('}'):
                        try:
                            next_parsed = json.loads(next_line)
                            next_val = str(list(next_parsed.values())[0])
                            if next_val in ('PASSED', 'FAILED'):
                                break
                        except (json.JSONDecodeError, IndexError):
                            pass
                    
                    # Skip separator lines and STDERR markers
                    if next_line == separator or next_line == dash_separator:
                        j += 1
                        continue
                    if next_line == stderr_begin or next_line == stderr_end:
                        j += 1
                        continue
                    
                    # Skip empty lines at the beginning
                    if not log_lines and not next_line:
                        j += 1
                        continue
                    
                    log_lines.append(lines[j].rstrip())
                    j += 1
                
                # Remove trailing empty lines
                while log_lines and not log_lines[-1].strip():
                    log_lines.pop()
                
                if log_lines:
                    log_text = '\n'.join(log_lines)
                    # Excel cell limit is 32767 characters
                    if len(log_text) > 32767:
                        log_text = log_text[:32764] + '...'
                    failed_logs[normalized_name] = log_text
                
                i = j  # Skip to where we stopped
                continue
            except (json.JSONDecodeError, IndexError):
                pass
        
        i += 1
    
    return failed_logs


def run_report(checklist_file, config_file, file_input):
    with open(file_input, "r", encoding="utf8") as file:
        line = file.read().splitlines()
    wb = load_workbook(filename=checklist_file)
    ws = wb.worksheets[0]
    with open(config_file, 'r', encoding="utf8") as j:
        contents = json.loads(j.read())
    index_audit = contents['data']

    mandatory_rows = set()
    for row_idx in range(1, ws.max_row + 1):
        cell_val = ws.cell(row=row_idx, column=5).value
        if cell_val and str(cell_val).strip().lower() == 'x':
            mandatory_rows.add(row_idx)

    mandatory_passed = 0
    mandatory_failed = 0
    optional_passed = 0
    optional_failed = 0

    failed_logs = extract_failed_logs(line)

    audit_results = {}
    for one_audit in line:
        try:
            parsed_data = json.loads(one_audit)
            result = str(list(parsed_data.keys())[0])
            value = str(list(parsed_data.values())[0])
            normalized_key = string_normalization(result)
            audit_results[normalized_key] = (value, one_audit)
        except Exception:
            continue

    for audit_pattern in index_audit:
        index_a = list(audit_pattern.keys())[0]
        index_pattern = string_normalization(list(audit_pattern.values())[0])

        entry = audit_results.get(index_pattern)
        if entry is None:
            continue

        status, raw_line = entry
        is_mandatory = int(index_a) in mandatory_rows if index_a.isdigit() else False

        if "FAILED" in raw_line:
            ws['D{}'.format(index_a)] = 'x'
            if index_pattern in failed_logs:
                ws['F{}'.format(index_a)] = failed_logs[index_pattern]
            if is_mandatory:
                mandatory_failed += 1
            else:
                optional_failed += 1
        else:
            ws['C{}'.format(index_a)] = 'x'
            if is_mandatory:
                mandatory_passed += 1
            else:
                optional_passed += 1

    # Find hostname and time dynamically in the file
    hostname = "unknown"
    time_generate = "unknown"
    
    for l in line:
        if l.strip().startswith("Hostname:"):
            hostname = l.replace("Hostname:", "").strip().split(" ")[0]
        elif l.strip().startswith("Audit Time:"):
            time_generate = l.replace("Audit Time:", "").strip().replace("-", "_").replace(":", "_").replace(" ", "-")
        # Also support old format with "Time:"
        elif l.strip().startswith("Time:") and time_generate == "unknown":
            time_generate = l.replace("Time:", "").strip().replace("-", "_").replace(":", "_").replace(" ", "-")
    
    # Create file name based strictly on hostname and time
    file_name = f'{hostname}_{time_generate}.xlsx'
    
    wb.save(filename=file_name)
    if check_file_exist(file_name):
        print(Fore.GREEN + "\nChecklist Done: {}".format(file_name) + Fore.RESET)
    else:
        print(Fore.RED + "Error creating Excel file!" + Fore.RESET)
    
    # Extract and return system info
    system_info = extract_system_info(line)
    
    # Add mandatory/optional statistics
    mandatory_total = mandatory_passed + mandatory_failed
    optional_total = optional_passed + optional_failed
    system_info['mandatory_passed'] = mandatory_passed
    system_info['mandatory_failed'] = mandatory_failed
    system_info['mandatory_total'] = mandatory_total
    system_info['mandatory_percentage'] = round((mandatory_passed / mandatory_total) * 100, 1) if mandatory_total > 0 else 0.0
    system_info['optional_passed'] = optional_passed
    system_info['optional_failed'] = optional_failed
    system_info['optional_total'] = optional_total
    system_info['optional_percentage'] = round((optional_passed / optional_total) * 100, 1) if optional_total > 0 else 0.0
    
    # Recalculate compliance stats based on config-matched criteria (consistent with Excel)
    # extract_system_info counts ALL PASSED/FAILED lines in raw text (including criteria
    # not in the checklist config), which can differ from the Excel result.
    # Use the actual config-matched counts for accurate web display.
    config_passed = mandatory_passed + optional_passed
    config_failed = mandatory_failed + optional_failed
    config_total = config_passed + config_failed
    if config_total > 0:
        system_info['passed_count'] = config_passed
        system_info['failed_count'] = config_failed
        system_info['total_checks'] = config_total
        system_info['compliance_percentage'] = round((config_passed / config_total) * 100, 1)
    
    return system_info


def run_generate_excel(checklist_file, config_file, filename):
    global current
    current = datetime.datetime.now()
    file_decrypt = os.path.basename(filename).replace(".enc", "")
    if check_file_exist(file_decrypt):
        os.remove(file_decrypt)
        decryption_write(filename)
    else:
        decryption_write(filename)
    if check_file_exist(file_decrypt):
        system_info = run_report(checklist_file, config_file, file_decrypt)
        return system_info
    return None