import os
import sys
import subprocess


# PATCH_FILE = "sts2019-kb5002678-fullfile-x64-glb.exe"
# OUTPUT_DIR = "src_feb"
PATCH_FILE = sys.argv[1]
OUTPUT_DIR = sys.argv[2]
# Decompress the file sts-x-none.msp from PATCH_FILE using 7-Zip
seven_zip_path = r"C:\Program Files\7-Zip\7z.exe"
ilspy_path = r"E:\project\ILSpy\ICSharpCode.ILSpyCmd\bin\x64\Release\net8.0\ilspycmd.exe"
patch_file_path = os.path.join(os.getcwd(), PATCH_FILE)
output_dir_path = os.path.join(os.getcwd(), OUTPUT_DIR)

# Ensure the output directory exists
if not os.path.exists(output_dir_path):
    os.makedirs(output_dir_path)
else:
    # Clear the output directory if it is not empty
    for root, dirs, files in os.walk(output_dir_path, topdown=False):
        for name in files:
            os.remove(os.path.join(root, name))
        for name in dirs:
            os.rmdir(os.path.join(root, name))

command = [
    seven_zip_path,
    'x',
    patch_file_path,
    'sts-x-none.msp',
    '-o' + output_dir_path,
    '-y',
    '-bd'
]

try:
    result = subprocess.run(command, check=True, capture_output=True, text=True)
    print("[+] Stage 1: success")
    sts_output = os.path.join(output_dir_path, "sts-x-none.msp")
    if not os.path.exists(sts_output):
        print("[-] sts-x-none.msp extracted failed")
        exit(0)
    print("[+] sts-x-none.msp extracted successfully")
    command = [
        seven_zip_path,
        'x',
        sts_output,
        'PATCH_CAB',
        '-o' + output_dir_path,
        '-y',
        '-bd'
    ]
    result = subprocess.run(command, check=True, capture_output=True, text=True)
    os.remove(sts_output)
    print("[+] Stage 2: success")
    patch_cab = os.path.join(output_dir_path, "PATCH_CAB")
    if not os.path.exists(patch_cab):
        print("[-] PATCH_CAB extracted failed")
        exit(0)
    print("[+] PATCH_CAB extracted successfully")
    dll_path = os.path.join(output_dir_path, "dll")
    if not os.path.exists(dll_path):
        os.makedirs(dll_path)
    command = [
        seven_zip_path,
        'x',
        patch_cab,
        '*.dll',
        '-o' + dll_path,
        '-y',
        '-bd'
    ]
    result = subprocess.run(command, check=True, capture_output=True, text=True)
    os.remove(patch_cab)
    print("[+] Stage 3: success")
    command = [
        ilspy_path,
        '-p',
        '--nested-directories',
        '-o',
        output_dir_path,
        dll_path
    ]
    print("[+] Start to decompile dll files")
    result = subprocess.run(command, check=True, capture_output=True, text=True)
    print("[+] Stage 4: success")
    for root, dirs, files in os.walk(output_dir_path):
        for file in files:
            if not file.endswith(".cs") or file == "AssemblyInfo.cs":
                os.remove(os.path.join(root, file))
    print("[+] Non-.cs files removed successfully")
    
    # print(result.stdout)
except subprocess.CalledProcessError as e:
    print("Command failed with error:")
    print(e.stderr)