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 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 test_ready_disconnected(self, mock_socket):
buffer = BufMock(Auth(self.challenge), Ready(0, "bub"))
mock_socket.recv = buffer.read
AgentHandler(mock_socket, ('127.0.0.1', '1234'), self.mock_server,
terminate_event=self.mock_terminate, context=self.ctx)
mock_socket.sendall.assert_called_with(
b'\x03\t\x10127.0.0.1-sV -Pn -p1-1000')
self.ctx.running.assert_not_called()
self.ctx.interrupted.assert_called_once()
def test_ready(self, mock_socket):
buffer = BufMock(Auth(self.challenge), Ready(0, "bub"),
struct.pack("<B", 0))
mock_socket.recv = buffer.read
AgentHandler(mock_socket, ('127.0.0.1', '1234'), self.mock_server,
terminate_event=self.mock_terminate, context=self.ctx)
mock_socket.sendall.assert_called_with(
b'\x03\t\x10127.0.0.1-sV -Pn -p1-1000')
self.ctx.running.assert_called_once()
self.ctx.running.assert_called_with("127.0.0.1:1234")
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_authentication(self, mock_urandom, mock_socket):
mock_urandom.return_value = self.challenge
# emulate the client!
hmac_hash = hmac.new(self.settings.secret_key, self.challenge,
'sha512')
digest = hmac_hash.hexdigest().encode("utf-8")
buffer = BufMock(Auth(digest))
mock_socket.recv = buffer.read
mock = MagicMock()
handler = AgentHandler(mock_socket, ('127.0.0.1', '1234'),
self.mock_server, terminate_event=self.mock_terminate,
context=self.ctx)
self.assertTrue(handler.authenticated)
def test_wait(self, mock_socket):
buffer = BufMock(Auth(self.challenge), Ready(0, "bub"),
Ready(0, "bub"), struct.pack("<B", 0))
mock_socket.recv = buffer.read
mock_data = MagicMock()
mock_data.side_effect = [None, ("127.0.0.1", "-sV -Pn -p1-1000")]
self.ctx.pop = mock_data
AgentHandler(mock_socket, ('127.0.0.1', '1234'), self.mock_server,
terminate_event=self.mock_terminate, context=self.ctx)
mock_socket.sendall.assert_called_with(
b'\x03\t\x10127.0.0.1-sV -Pn -p1-1000')
self.ctx.running.assert_called_once()
self.ctx.running.assert_called_with("127.0.0.1:1234")
def test_report_send(self, mock_socket):
expected_hash = "055a61499ea7c0d96332cf850f69ecb7295"
file = open(os.path.join(data_path, 'discovery-nonstandar.xml'),
'rb')
file.seek(0, os.SEEK_END)
report_msg = Report(file.tell(), "foobar.xml", expected_hash)
file.seek(0)
buffer = BufMock(Auth("fu").pack(), report_msg.pack(), file)
mock_socket.recv = buffer.read
report_mock = mock_open()
with patch('builtins.open', report_mock) as m:
self.ctx.get_report = report_mock
AgentHandler(mock_socket, ('127.0.0.1', '1234'), self.mock_server,
terminate_event=self.mock_terminate, context=self.ctx)
handle = m.return_value
offset = file.tell() - 996
file.seek(offset)
self.assertEqual(handle.write.call_count, 3)
handle.write.assert_called_with(file.read())
self.ctx.downloading.assert_called_with("127.0.0.1:1234")
self.ctx.completed.assert_called_with("127.0.0.1:1234")
file.close()