from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import sys
import os


def decrypt_file(input_file, key_hex):
    try:
        key = bytes.fromhex(key_hex)
    except ValueError:
        print("[-] Invalid AES key")
        return

    with open(input_file, "rb") as f:
        encrypted_data = f.read()

    cipher = AES.new(key, AES.MODE_ECB)
    dec_data = cipher.decrypt(pad(encrypted_data, 16))

    if b"INNER_MSTAR" not in dec_data:
        print("[-] AES key not working")
        return

    payload = None
    for offset in (64, 96):
        candidate = dec_data[offset:]
        if b"CHAI" in candidate or b"kbox" in candidate:
            payload = candidate
            break

    if payload is None:
        payload = dec_data[64:]

    if b"CHAI" in payload:
        out_data = payload
    elif b"kbox" in payload:
        out_data = payload[:128]
    else:
        out_data = payload[:32]

    base, _ = os.path.splitext(input_file)
    out_file = f"{base}_encrypted.dat"

    with open(out_file, "wb") as f:
        f.write(out_data)

    print(f"[+] Saved {out_file}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: python3 {sys.argv[0]} <file.dat>")
        sys.exit(1)

    aes_key = input("AES key (hex): ").strip()
    decrypt_file(sys.argv[1], aes_key)