from hashlib import pbkdf2_hmac
import os
import struct
import sys
import re

import pefile

def rva2offset(pe: pefile.PE, rva: int):
    for section in pe.sections:
        if section.contains_rva(rva):
            return section.get_offset_from_rva(rva)
    
    return -1

if len(sys.argv) < 2:
    print(f"Usage: python {os.path.basename(__file__)} <PATH TO EXE>")
    exit(1)

exe_path = sys.argv[1]

KEY_PASSWORD_RE = re.compile(rb"\?AVSystemInterface@projClient@@....(?P<offset>....)")
KEY_SALT_RE = re.compile(rb"\x50\x6A\x20\x6A(?P<iterations>.)\x6A\x10\x2B\xCA\x68(?P<offset>....)\x51\x55\xE8....")
SALT_PASSWORD_RE = re.compile(rb"\?AVDeflate@projClient@@\x00\x00\x00\x00....(?P<offset1>....)(?P<offset2>....)")
SALT_SALT_RE = re.compile(rb"\x52\x6A\x08\x6A(?P<iterations>.)\x6A\x10\x68(?P<offset>....)\x51\x53\xE8....")
IV_RE_1 = re.compile(rb"\xF3\x0F\x7E\x05(?P<offset>....)\x8B\x74\x24\x24\x6A\x01")
IV_RE_2 = re.compile(rb"\xE8....\xF3\x0F\x7E\x05(?P<offset>....)\x6A\x01")
ITER_COUNT_RE = re.compile(rb"\xC7\x86....(?P<count>....)\x0F\x8C....\x85\xED\x0F\x84....\x85\xDB\x0F\x84....")

with open(exe_path, "rb") as f:
    exe = f.read()

pe = pefile.PE(data=exe, fast_load=True)
base_address = pe.OPTIONAL_HEADER.ImageBase

if (pmatch := KEY_PASSWORD_RE.search(exe)) and (smatch := KEY_SALT_RE.search(exe)):
    poffset = rva2offset(pe, struct.unpack("<I", pmatch.group("offset"))[0] - base_address)
    soffset = rva2offset(pe, struct.unpack("<I", smatch.group("offset"))[0] - base_address)
    poffset_end = poffset + exe[poffset:].index(0)

    password = exe[poffset:poffset_end]
    salt = exe[soffset:soffset + 16]

    key = pbkdf2_hmac("sha1", password, salt, iterations=smatch.group("iterations")[0], dklen=32)
    key = bytes((x % 0x5E) + 0x21 for x in key)

    print(f"Key: {key.hex()}")

if (ivmatch := IV_RE_1.search(exe)):
    ivoffset = rva2offset(pe, struct.unpack("<I", ivmatch.group("offset"))[0] - base_address)
    print(f"IV: {exe[ivoffset:ivoffset + 16].hex()}")
elif (ivmatch := IV_RE_2.search(exe)):
    ivoffset = rva2offset(pe, struct.unpack("<I", ivmatch.group("offset"))[0] - base_address)
    print(f"IV: {exe[ivoffset:ivoffset + 16].hex()}")

if (pmatch := SALT_PASSWORD_RE.search(exe)) and (smatch := SALT_SALT_RE.search(exe)):
    poffset = rva2offset(pe, struct.unpack("<I", pmatch.group("offset1"))[0] - base_address)

    if poffset == -1:
        poffset = rva2offset(pe, struct.unpack("<I", pmatch.group("offset2"))[0] - base_address)

    soffset = rva2offset(pe, struct.unpack("<I", smatch.group("offset"))[0] - base_address)
    poffset_end = poffset + exe[poffset:].index(0)

    password = exe[poffset:poffset_end]
    salt = exe[soffset:soffset + 16]

    key = pbkdf2_hmac("sha1", password, salt, iterations=smatch.group("iterations")[0], dklen=8)
    print(f"Endpoint salt: {key.hex()}")

if (match := ITER_COUNT_RE.search(exe)):
    iter_count = struct.unpack("<I", match.group("count"))[0]
    
    print(f"Iterations: {iter_count}")