#!/usr/bin/env python

import re, os, sys
from hashlib import sha256
# This program uses an extended version of Swiftloke's brilliant MSProbe to assemble the payload
from MSProbe.assemble import asmMain

# match this many hexadigits
# must be corroborated within the script
depth = 6
# bytes per message
message_size = 1

shellcode_asm = 'shellcode.asm'
if len(sys.argv) > 1:
    shellcode_asm = sys.argv[1]

shellcode_out = f'{shellcode_asm}.tmp'

# Compile shellcode w/ msprobe
asmMain(shellcode_asm, shellcode_out, silent=True)

# Read compiled output
with open(shellcode_out) as file:
    shellcode = file.readline()
os.remove(shellcode_out)
shellcode_len = len(bytes.fromhex(shellcode));

# Print formatted payload as hex
print(f"6000{shellcode_len:x}{shellcode}")

def main():
    internal_sram_hash = input("Internal SRAM Hash: ").lower()

    target = re.sub(r'\s+', '', input("\nEncoded Target: ")).lower()

    # calculate LUT for given parameters
    all_hashes = {}
    for i in range (0, 1<<(8*message_size)):
        key = sha256(i.to_bytes(length=message_size, byteorder='big', signed=False)).hexdigest()[:depth]
        all_hashes[key] = i
    dedup = list(dict.fromkeys(all_hashes.keys()))
    #print(f"distinct values in all_hashes: {len(dedup)}\n")

    sram: bytearray = bytearray()

    # decode
    print("Decoded:")
    for loc in range (0, len(target), depth*message_size):
        key = target[loc:loc+depth]
        try:
            value = all_hashes[key]
            print(f"{value:02x}", end="")

        except KeyError:
            value = 0x00
        sram.append(value)

    # write to file
    with open("carfax.bin", 'wb') as output:
        output.write(sram)

    print("")
    # hexdump file
    os.system("xxd carfax.bin")

    # check hash against provided value
    while len(sram) < 0x1000:
        sram.append (0)
    print (f"\nHash:\nProvided: {internal_sram_hash}\n Decoded: {sha256(sram).hexdigest()}")

    if (sha256(sram).hexdigest().lower() == internal_sram_hash.lower()):
        print("Hash match!")

if __name__ == "__main__":
    main()