def test_server(self, mock_open):
with patch('os.path.isfile') as misfile:
misfile.return_value = False
server = DScanServer((self.settings.host, self.settings.port),
AgentHandler, options=self.settings)
server_thread = threading.Thread(target=server.serve_forever)
# Exit the server thread when the main thread terminates
server_thread.daemon = True
server_thread.start()
log.info(f"Server loop running in thread:{server_thread.name}")
context = ssl.create_default_context()
context.check_hostname = False
context.load_verify_locations(self.settings.sslcert)
s = context.wrap_socket(socket(AF_INET, SOCK_STREAM),
server_side=False, server_hostname="dscan")
s.connect(('127.0.0.1', 9011))
opr = Structure.create(s)
hmac_hash = hmac.new(self.settings.secret_key, opr.data,
'sha512')
digest = hmac_hash.hexdigest().encode("utf-8")
s.sendall(Auth(digest).pack())
s.close()
server.shutdown()
def test_result_status_pack_unpack(self):
expected = ExitStatus(Status.FINISHED)
mock_sock = self.build_mock(expected)
with patch('socket.socket', new=mock_sock) as mock_socket:
result = Structure.create(sock=mock_socket)
self.assertEqual(expected.op_code, result.op_code)
self.assertEqual(expected.status, result.status)
def test_ready_pack_unpack(self):
expected = Ready(0, 'agentx')
mock_sock = self.build_mock(expected)
with patch('socket.socket', new=mock_sock) as mock_socket:
result = Structure.create(sock=mock_socket)
self.assertEqual(expected.uid, result.uid)
self.assertEqual(expected.alias, result.alias)
def do_auth(self):
"""
Handles the agent's authentication.
"""
log.info(f"{self.client_address} initialized Authentication")
challenge = os.urandom(128)
self.request.sendall(Auth(challenge).pack())
# wait for the client response
self.msg = Structure.create(self.request)
if not self.msg:
# something's wrong with the message!
self.send_status(Status.FAILED)
return
hmac_hash = hmac.new(self.server.secret_key, challenge, 'sha512')
digest = hmac_hash.hexdigest().encode("utf-8")
if hmac.compare_digest(digest, self.msg.data):
log.info("Authenticated Successfully")
self.authenticated = True
self.send_status(Status.SUCCESS)
else:
self.send_status(Status.UNAUTHORIZED)
self.request.close()
def dispatcher(self):
"""
Command dispatcher all logic to decode and dispatch the call.
"""
self.msg = Structure.create(self.request)
if not self.msg:
self.connected = False
log.info("Disconnected!")
# mark any running task as interrupted
# so that other agent can take it later
self.ctx.interrupted(self.agent)
return
command_name = f"do_{self.msg.op_code.name.lower()}"
if not hasattr(self, command_name):
self.send_status(Status.FAILED) # invalid command
return
command = getattr(self, command_name)
# the only command authorized for unauthenticated agents
if command_name != "do_AUTH" and not self.authenticated:
self.send_status(Status.UNAUTHORIZED)
self.connected = False
self.request.close()
return
# call the command !
command()
def test_report_pack_unpack(self):
digest = hashlib.sha512(b"pickabu").hexdigest()
expected = Report(self.getSize(), "fu.xml", digest)
mock_sock = self.build_mock(expected)
with patch('socket.socket', new=mock_sock) as mock_socket:
result = Structure.create(sock=mock_socket)
self.assertEqual(expected.filesize, result.filesize)
self.assertEqual(expected.filename, result.filename)
self.assertEqual(expected.filehash, result.filehash)
def test_auth_pack_unpack(self):
digest = hashlib.sha512(b"pickabu").hexdigest()
expected = Auth(digest)
mock_sock = self.build_mock(expected)
with patch('socket.socket', new=mock_sock) as mock_socket:
struct_size: int = struct.calcsize("<B")
result = Structure.create(sock=mock_socket)
self.assertEqual(expected.data, result.data)
self.assertEqual(digest.encode("ascii"), result.data)
self.assertEqual(Operations.AUTH, result.op_code)
def test_command_pack_unpack(self):
expected = Command("127.0.0.1", "-sV -Pn -p1-1000")
mock_sock = self.build_mock(expected)
with patch('socket.socket', new=mock_sock) as mock_socket:
result = Structure.create(sock=mock_socket)
self.assertEqual(expected.target, result.target)