Update agent server parse

app/objects/c_agent.py

@@ -107,8 +107,7 @@ def __init__(self, sleep_min=30, sleep_max=60, watchdog=0, platform='unknown', s
         self.group = group
         self.architecture = architecture
         self.platform = platform.lower()
-        url = urlparse(server)
-        self.server = '%s://%s:%s' % (url.scheme, url.hostname, url.port)
+        self.server = self.parse_endpoint(server)
         self.location = location
         self.pid = pid
         self.ppid = ppid
@@ -132,11 +131,7 @@ def __init__(self, sleep_min=30, sleep_max=60, watchdog=0, platform='unknown', s
         self.available_contacts = available_contacts if available_contacts else [self.contact]
         self.pending_contact = pending_contact if pending_contact else contact
         self.host_ip_addrs = host_ip_addrs if host_ip_addrs else []
-        if upstream_dest:
-            upstream_url = urlparse(upstream_dest)
-            self.upstream_dest = '%s://%s:%s' % (upstream_url.scheme, upstream_url.hostname, upstream_url.port)
-        else:
-            self.upstream_dest = self.server
+        self.upstream_dest = self.parse_endpoint(upstream_dest) if upstream_dest else self.server
         self._executor_change_to_assign = None
         self.log = self.create_logger('agent')
 
@@ -195,7 +190,8 @@ async def heartbeat_modification(self, **kwargs):
             self.last_trusted_seen = now
         self.update('pid', kwargs.get('pid'))
         self.update('ppid', kwargs.get('ppid'))
-        self.update('server', kwargs.get('server'))
+        new_server = kwargs.get('server')
+        self.update('server', self.parse_endpoint(new_server) if new_server else None)
         self.update('exe_name', kwargs.get('exe_name'))
         self.update('location', kwargs.get('location'))
         self.update('privilege', kwargs.get('privilege'))
@@ -331,6 +327,20 @@ def assign_pending_executor_change(self):
         self._executor_change_to_assign = None
         return executor_change
 
+    @staticmethod
+    def parse_endpoint(server):
+        if re.search(r'^[^:/\s]+://(?=([^:\s]+))\1(:\d{1,5})?/*$', server):
+            # E.g. http://127.0.0.1:8888
+            url = urlparse(server)
+            parsed = '%s://%s' % (url.scheme, url.hostname)
+            return parsed + ':%s' % url.port if url.port else parsed
+        elif re.search(r'^[^\s:]+:\d{1,5}/*$', server):
+            # E.g. 127.0.0.1:7010 or 127.0.0.1:7011/
+            return server.rstrip('/')
+        return 'unknown'
+
+    """ PRIVATE """
+
     def _replace_payload_data(self, decoded_cmd, file_svc):
         for uuid in re.findall(self.RESERVED['payload'], decoded_cmd):
             if self.is_uuid4(uuid):

tests/api/v2/handlers/test_operations_api.py

@@ -352,7 +352,7 @@ async def test_create_potential_link_with_globals(self, api_v2_client, api_cooki
             assert result['paw'] == payload['paw']
             assert result['id']
             assert result['ability']['name'] == 'Manual Command'
-            assert result['command'] == "://None:None 123"
+            assert result['command'] == "unknown 123"
 
     async def test_create_potential_link(self, api_v2_client, api_cookies, mocker, async_return):
         with mocker.patch('app.objects.c_operation.Operation.apply') as mock_apply:

tests/objects/test_agent.py

@@ -126,6 +126,91 @@ def test_heartbeat_modification_during_pending_executor_removal(self, event_loop
         event_loop.run_until_complete(agent.heartbeat_modification(executors=original_executors))
         assert agent.executors == ['cmd']
 
+    def test_server_value_with_protocol_ip_port(self):
+        test = 'http://127.0.0.1:8888'
+        test_slash = 'http://127.0.0.1:8888/'
+        want = 'http://127.0.0.1:8888'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_with_protocol_ip_without_port(self):
+        test = 'http://127.0.0.1'
+        test_slash = 'http://127.0.0.1/'
+        want = 'http://127.0.0.1'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_with_protocol_domain_port(self):
+        test = 'http://mydomain.tld:8888'
+        test_slash = 'http://mydomain.tld:8888/'
+        want = 'http://mydomain.tld:8888'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_with_protocol_domain_without_port(self):
+        test = 'http://mydomain.tld'
+        test_slash = 'http://mydomain.tld/'
+        want = 'http://mydomain.tld'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_without_protocol_with_ip_port(self):
+        test = '127.0.0.1:7010'
+        test_slash = '127.0.0.1:7010/'
+        want = '127.0.0.1:7010'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_without_protocol_with_domain_port(self):
+        test = 'mydomain.tld:7010'
+        test_slash = 'mydomain.tld:7010/'
+        want = 'mydomain.tld:7010'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_with_protocol_localhost_port(self):
+        test = 'http://localhost:8888'
+        test_slash = 'http://localhost:8888/'
+        want = 'http://localhost:8888'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_with_protocol_localhost_without_port(self):
+        test = 'http://localhost'
+        test_slash = 'http://localhost/'
+        want = 'http://localhost'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_without_protocol_localhost_port(self):
+        test = 'localhost:8888'
+        test_slash = 'localhost:8888/'
+        want = 'localhost:8888'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_server_value_without_protocol_localhost_without_port(self):
+        test = 'localhost'
+        test_slash = 'localhost/'
+        want = 'unknown'
+        assert Agent.parse_endpoint(test) == want
+        assert Agent.parse_endpoint(test_slash) == want
+
+    def test_heartbeat_modification_new_server(self, loop):
+        agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['cmd', 'test'], platform='windows',
+                      server='unknown')
+        test = 'mydomain.tld:7010'
+        want = 'mydomain.tld:7010'
+        loop.run_until_complete(agent.heartbeat_modification(server=test))
+        assert agent.server == want
+
+    def test_heartbeat_modification_no_server_provided(self, loop):
+        want = 'mydomain.tld:7010'
+        agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['cmd', 'test'], platform='windows',
+                      server=want)
+        loop.run_until_complete(agent.heartbeat_modification())
+        assert agent.server == want
+
     def test_store_new_agent(self, data_svc):
         agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['cmd', 'test'], platform='windows')
         stored_agent = agent.store(data_svc.ram)

tests/services/test_rest_svc.py

@@ -82,13 +82,13 @@ def test_delete_operation(self, event_loop, rest_svc, data_svc):
                               'host_group': [{'trusted': True, 'architecture': 'unknown', 'watchdog': 0,
                                               'contact': 'unknown', 'username': 'unknown', 'links': [], 'sleep_max': 8,
                                               'exe_name': 'unknown', 'executors': ['pwsh', 'psh'], 'ppid': 0,
-                                              'sleep_min': 2, 'server': '://None:None', 'platform': 'windows',
+                                              'sleep_min': 2, 'server': 'unknown', 'platform': 'windows',
                                               'host': 'unknown', 'paw': '123', 'pid': 0,
                                               'display_name': 'unknown$unknown', 'group': 'red', 'location': 'unknown',
                                               'privilege': 'User', 'proxy_receivers': {}, 'proxy_chain': [],
                                               'origin_link_id': '', 'deadman_enabled': False,
                                               'available_contacts': ['unknown'], 'pending_contact': 'unknown',
-                                              'host_ip_addrs': [], 'upstream_dest': '://None:None'}],
+                                              'host_ip_addrs': [], 'upstream_dest': 'unknown'}],
                               'visibility': 50, 'autonomous': 1, 'chain': [], 'auto_close': False,
                               'obfuscator': 'plain-text', 'use_learning_parsers': False,
                               'group': '',
@@ -158,12 +158,12 @@ def test_create_operation(self, event_loop, rest_svc, data_svc):
                 'host_group': [
                     {'trusted': True, 'architecture': 'unknown', 'watchdog': 0, 'contact': 'unknown',
                      'username': 'unknown', 'links': [], 'sleep_max': 8, 'exe_name': 'unknown',
-                     'executors': ['pwsh', 'psh'], 'ppid': 0, 'sleep_min': 2, 'server': '://None:None',
+                     'executors': ['pwsh', 'psh'], 'ppid': 0, 'sleep_min': 2, 'server': 'unknown',
                      'platform': 'windows', 'host': 'unknown', 'paw': '123', 'pid': 0,
                      'display_name': 'unknown$unknown', 'group': 'red', 'location': 'unknown', 'privilege': 'User',
                      'proxy_receivers': {}, 'proxy_chain': [], 'origin_link_id': '',
                      'deadman_enabled': False, 'available_contacts': ['unknown'], 'pending_contact': 'unknown',
-                     'host_ip_addrs': [], 'upstream_dest': '://None:None'}],
+                     'host_ip_addrs': [], 'upstream_dest': 'unknown'}],
                 'visibility': 50, 'autonomous': 1, 'chain': [], 'auto_close': False, 'objective': '',
                 'obfuscator': 'plain-text', 'use_learning_parsers': False}
         internal_rest_svc = rest_svc(event_loop)