"""
Windows Audit Hardening Tool - v2.0
Usage: AuditTool.exe -p audit_cis_windows.ps1.enc
       AuditTool.exe                  (uses bundled .enc file)

Build:  python -m PyInstaller --onefile --console -n AuditTool ^
            --add-data "audit_cis_windows.ps1.enc;." ^
            windows_audit.py

Environment: Python 3.8+ (3.8 for Windows Server 2012 compatibility)
"""

import argparse
import base64
import os
import socket
import subprocess
import sys
import tempfile
import traceback

from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.PublicKey import RSA
from Crypto.Util.Padding import unpad
from colorama import Fore, init as colorama_init
from unidecode import unidecode


# ----------------------------------------------------------
# Globals
# ----------------------------------------------------------
colorama_init(strip=False, autoreset=True)

HOSTNAME = socket.gethostname()
try:
    IP_ADDR = socket.gethostbyname(HOSTNAME)
except Exception:
    IP_ADDR = "127.0.0.1"

AES_KEY_B64 = "fb8MH7yIuVUeCp5W4S9cKFtsHaxXIP/zvkO36wNNcO4="
BUNDLED_ENC = "audit_cis_windows.ps1.enc"
PUBLIC_KEY_FILE = "public_key.pem"

BANNER = r"""
          _    _ _____ _____ _______            _    _          _____  _____  ______ _   _ _____ _   _  _____ 
     /\  | |  | |  __ \_   _|__   __|          | |  | |   /\   |  __ \|  __ \|  ____| \ | |_   _| \ | |/ ____|
    /  \ | |  | | |  | || |    | |     ______  | |__| |  /  \  | |__) | |  | | |__  |  \| | | | |  \| | |  __ 
   / /\ \| |  | | |  | || |    | |    |______| |  __  | / /\ \ |  _  /| |  | |  __| | . ` | | | | . ` | | |_ |
  / ____ \ |__| | |__| || |_   | |             | |  | |/ ____ \| | \ \| |__| | |____| |\  |_| |_| |\  | |__| |
 /_/    \_\____/|_____/_____|  |_|             |_|  |_/_/    \_\_|  \_\_____/|______|_| \_|_____|_| \_|\_____|
"""


# ----------------------------------------------------------
# Path resolution (supports PyInstaller bundle)
# ----------------------------------------------------------
def resource_path(relative_path):
    """Get absolute path to resource, works for dev and PyInstaller."""
    if getattr(sys, "frozen", False):
        base = sys._MEIPASS
    else:
        base = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(base, relative_path)


def resolve_enc_path(user_path=None):
    """
    Resolve the .enc file path.
    Priority: user-supplied arg > bundled file > same-dir file.
    Returns absolute path or None.
    """
    if user_path:
        if os.path.isfile(user_path):
            return os.path.abspath(user_path)
        print(Fore.YELLOW + "[WARNING] Provided path not found: {}".format(user_path) + Fore.RESET)

    # Try PyInstaller bundled location
    bundled = resource_path(BUNDLED_ENC)
    if os.path.isfile(bundled):
        return bundled

    # Try same directory
    same_dir = os.path.join(os.getcwd(), BUNDLED_ENC)
    if os.path.isfile(same_dir):
        return same_dir

    return None


def resolve_public_key():
    """Find public_key.pem - bundled, same dir, or current dir."""
    for loc in [
        resource_path(PUBLIC_KEY_FILE),
        os.path.join(os.getcwd(), PUBLIC_KEY_FILE),
    ]:
        if os.path.isfile(loc):
            return loc
    return None


# ----------------------------------------------------------
# Crypto helpers
# ----------------------------------------------------------
def decrypt_aes(file_path, key_b64):
    """Decrypt an AES-ECB (PKCS7 padded) encrypted file."""
    key = base64.b64decode(key_b64)
    with open(file_path, "rb") as fh:
        ct = fh.read()
    cipher = AES.new(key, AES.MODE_ECB)
    padded = cipher.decrypt(ct)
    data = unpad(padded, AES.block_size, style="pkcs7")
    return data.decode("utf-8")


def encrypt_output_rsa(content):
    """Encrypt output with RSA-OAEP; fall back to plaintext if no key."""
    pk_path = resolve_public_key()
    if pk_path is None:
        return "PLAINTEXT:" + content
    try:
        pub_key = RSA.import_key(open(pk_path, "rb").read())
        cipher = PKCS1_OAEP.new(pub_key)
        raw = content.encode("utf-8")
        block_size = 190
        ciphertext = b""
        for i in range(0, len(raw), block_size):
            ciphertext += cipher.encrypt(raw[i : i + block_size])
        return base64.b64encode(ciphertext).decode("utf-8")
    except Exception as e:
        print(Fore.YELLOW + "[WARNING] RSA encrypt failed: {}".format(e) + Fore.RESET)
        return "PLAINTEXT:" + content


# ----------------------------------------------------------
# PowerShell runner
# ----------------------------------------------------------
def run_powershell_script(script_content, timeout_sec=600):
    """
    Write script to temp .ps1 file, execute via 'powershell -File',
    return (stdout, stderr, returncode).
    """
    tmp_path = None
    try:
        fd, tmp_path = tempfile.mkstemp(suffix=".ps1", prefix="audit_")
        with os.fdopen(fd, "w", encoding="utf-8") as fh:
            fh.write(script_content)

        proc = subprocess.run(
            [
                "powershell.exe",
                "-NoProfile",
                "-ExecutionPolicy", "Bypass",
                "-File", tmp_path,
            ],
            capture_output=True,
            text=True,
            timeout=timeout_sec,
            encoding="utf-8",
            errors="replace",
        )
        return proc.stdout, proc.stderr, proc.returncode
    except subprocess.TimeoutExpired:
        return "", "PowerShell execution timed out ({}s)".format(timeout_sec), -1
    except Exception:
        return "", traceback.format_exc(), -1
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.remove(tmp_path)
            except Exception:
                pass


# ----------------------------------------------------------
# Output parsing
# ----------------------------------------------------------
def extract_audit_info(output):
    """Extract hostname + timestamp from audit output lines."""
    hostname = HOSTNAME
    timestamp = "unknown"
    for line in output.splitlines():
        if line.startswith("Hostname:"):
            val = line.replace("Hostname:", "").strip()
            hostname = val if val else HOSTNAME
        elif line.startswith("Time:"):
            val = line.replace("Time:", "").strip()
            timestamp = val.replace("-", "_").replace(":", "_").replace(" ", "-") if val else "unknown"
    return hostname, timestamp


def count_pass_fail(output):
    """Count PASSED/FAILED from JSON lines."""
    passed = failed = 0
    for line in output.splitlines():
        line = line.strip()
        if line.startswith("{") and line.endswith("}"):
            if '"PASSED"' in line or "'PASSED'" in line:
                passed += 1
            elif '"FAILED"' in line or "'FAILED'" in line:
                failed += 1
    return passed, failed


# ----------------------------------------------------------
# Main audit flow
# ----------------------------------------------------------
def run_audit(enc_path=None):
    """Decrypt .enc, run PowerShell audit, encrypt+save results."""
    print(Fore.BLUE + BANNER + Fore.RESET)
    print(Fore.RED + "Running Audit Hardening v2.0" + Fore.RESET)
    print("  Hostname : {}".format(HOSTNAME))
    print("  IP       : {}".format(IP_ADDR))
    print("-" * 70)

    # Step 0 – Resolve input file
    enc_path = resolve_enc_path(enc_path)
    if enc_path is None:
        print(Fore.RED + "ERROR: No .enc file found.")
        print("  Provide path:  AuditTool.exe -p audit_cis_windows.ps1.enc")
        print("  Or place {} in the same directory.".format(BUNDLED_ENC) + Fore.RESET)
        return

    print("  Source   : {}".format(enc_path))

    # Step 1 – Decrypt
    print("  [1/4] Decrypting...")
    try:
        ps_content = decrypt_aes(enc_path, AES_KEY_B64)
    except Exception as e:
        print(Fore.RED + "  FAILED: {}".format(e) + Fore.RESET)
        return
    print("  [1/4] OK - {} bytes".format(len(ps_content)))

    # Step 2 – Execute PowerShell
    print("  [2/4] Running PowerShell...")
    stdout, stderr, rc = run_powershell_script(ps_content)
    if rc != 0:
        print(Fore.YELLOW + "  [2/4] PowerShell rc={}: {}".format(rc, (stderr or "")[:200]) + Fore.RESET)
    else:
        print("  [2/4] OK - {} output lines".format(len(stdout.splitlines())))

    # Step 3 – Combine
    full_output = stdout
    if stderr:
        full_output += "\n[STDERR]\n" + stderr
    if not full_output.strip():
        print(Fore.RED + "ERROR: No output from PowerShell" + Fore.RESET)
        return

    passed, failed = count_pass_fail(full_output)
    print("  [3/4] Audit: {} PASSED / {} FAILED".format(passed, failed))

    # Step 4 – Encrypt & save
    hostname_out, timestamp = extract_audit_info(stdout)
    out_filename = "{}_{}.txt.enc".format(hostname_out, timestamp)
    print("  [4/4] Encrypting -> {}".format(out_filename))

    encrypted = encrypt_output_rsa(full_output)
    with open(out_filename, "w", encoding="utf-8") as fh:
        fh.write(encrypted)

    if os.path.isfile(out_filename):
        print(Fore.GREEN + "=" * 70 + Fore.RESET)
        print(Fore.GREEN + "  SUCCESS: {}  |  {} PASS / {} FAIL".format(
            out_filename, passed, failed) + Fore.RESET)
        print(Fore.GREEN + "=" * 70 + Fore.RESET)
    else:
        print(Fore.RED + "  FAILED: Could not write output file" + Fore.RESET)


# ----------------------------------------------------------
# CLI
# ----------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Audit Hardening v2.0")
    parser.add_argument(
        "-p", "--path",
        help="Path to encrypted .ps1.enc file (optional if bundled or in same dir)",
        default=None,
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = main()
    run_audit(args.path)