attacker.py

import threading
import db_utils
import settings
import requests
from pwn import *
from submitter import CustomSubmitter


class Attacker(threading.Thread):
    def __init__(self, stngs: settings.Settings,
                 opponent: str, service: str, exploit, caller_globals,
                 exploit_type: str, queued_submitter):
        # TODO Type hint? -> builtins.function, need to install
        super().__init__()
        self.settings = stngs
        self.opponent = opponent
        self.service = service
        self.exploit = exploit
        self.stop = False
        self.killed = False
        self.errors = 0
        self.last_flag_time = time.time()
        self.elapsed_flag_time = self.last_flag_time
        self.stop_log = False
        self.type = exploit_type
        globals().update(caller_globals)
        self.queued_submitter = queued_submitter

    def get_valid_flags(self, flags: set) -> set:
        new_flags = set()
        for flag in flags:
            connection = db_utils.connect(self.settings.db)
            if not db_utils.check_if_a_flag_is_present_in_db(connection, flag):
                new_flags.add(flag)
            connection.close()
        return new_flags

    def submit_flags(self, flags: set, submitter: CustomSubmitter):
        print("[+] Exploit {} against {} about to submit flags: {}".format(self.service, self.opponent, flags))
        for flag in flags:
            connection = db_utils.connect(self.settings.db)
            self.__submit_flag(flag, submitter)
            db_utils.insert_flag_in_db(connection, flag,
                                       self.opponent, self.service)
            connection.close()

    def __submit_flag(self, flag: str, submitter: CustomSubmitter):
        data = {'team_token': self.settings.submit_token, 'flag': flag}
        try:
            if self.settings.submit_protocol == "POST":
                r = requests.post(self.settings.submit_url, data=data)
            elif self.settings.submit_protocol == "GET":
                r = requests.get(self.settings.submit_url, data=data)
            elif self.settings.submit_protocol == "TCP":
                url, port = self.settings.submit_url.split(":")
                conn = remote(url, port)
                conn.sendline(flag)
                conn.close()
            elif self.settings.submit_protocol == "CUSTOM":
                submitter.submit(flag, self.settings.submit_url, self.settings.submit_token)
                # submitter.submitter['submitter'](flag, self.settings.submit_url, self.settings.submit_token)
            elif self.settings.submit_protocol == 'CUSTOM-QUEUED':
                self.queued_submitter.enqueue_flag(flag)
        except requests.RequestException as e:
            msg = f"[+] An exception occurred while submitting a flag with requests:\n{e}"
            print(msg)
            raise Exception(msg)
        except PwnlibException as e:
            msg = f"[+] An exception occurred while submitting a flag with pwn:\n{e}"
            print(msg)
            raise Exception(msg)
        except Exception as e:
            msg = f"[+] An unknown exception occurred while submitting a flag:\n{e}"
            print(msg)
            raise Exception(msg)
        else:
            print(f"[+] Flag {flag} submitted")

    def kill(self):
        self.killed = True

    def switch_state(self):
        self.stop = not self.stop

    def is_active(self):
        return 1 if not self.stop else 0

    def log_stopping(self, cause):
        if not self.stop_log:
            msg = "[+] Exploit {} against {} stopping because of {}".format(self.service, self.opponent, cause)
            print(msg)

    def run(self):
        # auto-stop if it has error_threshold consecutive errors
        error_threshold = self.settings.thread_error_th
        submitter = CustomSubmitter(self.settings, globals())   # it is a singleton
        is_yield, generator = (self.type == 'yield'), None
        if is_yield:
            generator = self.exploit(self.opponent)
        while True:
            # auto-stop if it isn't finding new flags after some time
            low_performance = (self.elapsed_flag_time - self.last_flag_time) >= (2 * self.settings.tick)
            if (self.errors >= error_threshold) or low_performance:
                self.log_stopping("low performance" if low_performance else "too many errors")
                self.stop = True
                self.stop_log = True
            if self.killed:
                self.stop = True
                self.log_stopping("killing")
                self.stop_log = True
                break
            if not self.stop:
                self.stop_log = False
                print(f"[+] Now attacking target: {self.opponent}, service: {self.service}")
                try:
                    if is_yield:
                        flags = set(next(generator))
                    else:
                        flags = set(self.exploit(self.opponent))
                    new_flags = self.get_valid_flags(flags)
                    if len(new_flags) > 0:
                        self.last_flag_time = time.time()
                        self.elapsed_flag_time = self.last_flag_time
                    else:
                        # elapsed time tends to be equal to last flag time if new flags are always found
                        self.elapsed_flag_time = time.time()
                    self.submit_flags(new_flags, submitter)
                    if self.errors > 0:
                        self.errors -= 1
                except:
                    # TODO: log exceptions from here to file
                    self.errors += 1
            else:
                time.sleep(1)   # reduce CPU usage from stopped threads