Merge pull request #2713 from mitre/VIRTS-4153

[VIRTS-4153] ]Update server to account for exit_code from agent

app/objects/secondclass/c_result.py

@@ -7,6 +7,7 @@ class ResultSchema(ma.Schema):
     id = ma.fields.String()
     output = ma.fields.String()
     stderr = ma.fields.String()
+    exit_code = ma.fields.String()
     pid = ma.fields.String()
     status = ma.fields.String()
     agent_reported_time = ma.fields.DateTime(format=BaseObject.TIME_FORMAT, missing=None)
@@ -26,11 +27,12 @@ class Result(BaseObject):
 
     schema = ResultSchema()
 
-    def __init__(self, id, output, stderr="", pid=0, status=0, agent_reported_time=None):
+    def __init__(self, id, output, stderr='', exit_code='', pid=0, status=0, agent_reported_time=None):
         super().__init__()
         self.id = id
         self.output = output
         self.stderr = stderr
+        self.exit_code = exit_code
         self.pid = pid
         self.status = status
         self.agent_reported_time = agent_reported_time

app/service/contact_svc.py

@@ -126,7 +126,8 @@ async def _save(self, result):
                     result.output = await self._postprocess_link_result(result.output, link)
                     command_results = json.dumps(dict(
                         stdout=self.decode_bytes(result.output, strip_newlines=False),
-                        stderr=self.decode_bytes(result.stderr, strip_newlines=False)))
+                        stderr=self.decode_bytes(result.stderr, strip_newlines=False),
+                        exit_code=result.exit_code))
                     encoded_command_results = self.encode_string(command_results)
                     self.get_service('file_svc').write_result_file(result.id, encoded_command_results)
                     operation = await self.get_service('app_svc').find_op_with_link(result.id)
@@ -145,7 +146,8 @@ async def _save(self, result):
             else:
                 command_results = json.dumps(dict(
                     stdout=self.decode_bytes(result.output, strip_newlines=False),
-                    stderr=self.decode_bytes(result.stderr, strip_newlines=False)))
+                    stderr=self.decode_bytes(result.stderr, strip_newlines=False),
+                    exit_code=result.exit_code))
                 encoded_command_results = self.encode_string(command_results)
                 self.get_service('file_svc').write_result_file(result.id, encoded_command_results)
         except Exception as e:

app/service/file_svc.py

@@ -126,10 +126,10 @@ def read_result_file(self, link_id, location='data/results'):
             return decoded_buf
         except json.JSONDecodeError:
             results = json.dumps(dict(
-                stdout=self.decode_bytes(decoded_buf, strip_newlines=False), stderr=''))
+                stdout=self.decode_bytes(decoded_buf, strip_newlines=False), stderr='', exit_code=''))
             return self.encode_string(str(results))
         except binascii.Error:
-            results = json.dumps(dict(stdout=decoded_buf, stderr=''))
+            results = json.dumps(dict(stdout=decoded_buf, stderr='', exit_code=''))
             return self.encode_string(str(results))
 
     def write_result_file(self, link_id, output, location='data/results'):

templates/operations.html

@@ -427,13 +427,14 @@ <h2>Operations</h2>
                                 </div>
                             </template>
                             <div>
-                                <p class="mt-3 mb-2" x-text="`Standard Output: ${selectedLinkResults.stdout ? '' : 'Nothing to show'}`"></p>
+                                <p class="mt-3 mb-2" x-text="`Exit Code: ${selectedLinkResults && selectedLinkResults.exit_code ? selectedLinkResults.exit_code : 'Nothing to show'}`"></p>
+                                <p class="mt-3 mb-2" x-text="`Standard Output: ${selectedLinkResults && selectedLinkResults.stdout ? '' : 'Nothing to show'}`"></p>
                                 <template x-if="selectedLinkResults != null && selectedLinkResults.stdout !== ''">
                                     <pre class="has-text-left white-space-pre-line" x-text="selectedLinkResults.stdout"></pre>
                                 </template>
                             </div>
                             <div>
-                                <p class="mt-3 mb-2" x-text="`Standard Error: ${selectedLinkResults.stderr ? '' : 'Nothing to show'}`"></p>
+                                <p class="mt-3 mb-2" x-text="`Standard Error: ${selectedLinkResults && selectedLinkResults.stderr ? '' : 'Nothing to show'}`"></p>
                                 <template x-if="selectedLinkResults != null && selectedLinkResults.stderr !== ''">
                                     <pre class="has-text-left white-space-pre-line has-text-danger" x-text="selectedLinkResults.stderr"></pre>
                                 </template>

tests/services/test_contact_svc.py

@@ -45,12 +45,14 @@ class TestContactSvc:
     async def test_save_ability_hooks(self, setup_contact_service, contact_svc, event_svc):
         test_string = b'test_string'
         err_string = b'err_string'
+        test_exit_code = "-1"
         link = setup_contact_service
         rest_svc = RestService()
         result = dict(
             id=link.id,
             output=str(base64.b64encode(base64.b64encode(test_string)), 'utf-8'),
             stderr=str(base64.b64encode(err_string), 'utf-8'),
+            exit_code=test_exit_code,
             pid=0,
             status=0
         )
@@ -60,6 +62,7 @@ async def test_save_ability_hooks(self, setup_contact_service, contact_svc, even
 
         assert result_dict['stdout'] == test_string.decode()
         assert result_dict['stderr'] == err_string.decode()
+        assert result_dict['exit_code'] == test_exit_code
 
         # cleanup test
         try:
@@ -70,11 +73,13 @@ async def test_save_ability_hooks(self, setup_contact_service, contact_svc, even
     async def test_save_ability_hooks_with_no_link(self, setup_contact_service, contact_svc, event_svc, file_svc):
         test_string = b'test_string'
         err_string = b'err_string'
+        test_exit_code = "0"
         # Send version with link for comparison
         result = dict(
             id="12345",
             output=str(base64.b64encode(test_string), 'utf-8'),
             stderr=str(base64.b64encode(err_string), 'utf-8'),
+            exit_code=test_exit_code,
             pid=0,
             status=0
         )
@@ -85,6 +90,7 @@ async def test_save_ability_hooks_with_no_link(self, setup_contact_service, cont
         result_dict = json.loads(base64.b64decode(result))
         assert result_dict['stdout'] == test_string.decode()
         assert result_dict['stderr'] == err_string.decode()
+        assert result_dict['exit_code'] == test_exit_code
 
         # cleanup test
         try:

tests/services/test_file_svc.py

@@ -54,10 +54,11 @@ def test_read_write_result_file(self, tmpdir, file_svc):
         link_id = '12345'
         output = 'output testing unit'
         error = 'error testing unit'
-        output_encoded = str(b64encode(json.dumps(dict(stdout=output, stderr=error)).encode()), 'utf-8')
+        test_exit_code = '0'
+        output_encoded = str(b64encode(json.dumps(dict(stdout=output, stderr=error, exit_code=test_exit_code)).encode()), 'utf-8')
         file_svc.write_result_file(link_id=link_id, output=output_encoded, location=tmpdir)
 
-        expected_output = dict(stdout=output, stderr=error)
+        expected_output = dict(stdout=output, stderr=error, exit_code=test_exit_code)
         output_data = file_svc.read_result_file(link_id=link_id, location=tmpdir)
         decoded_output_data = json.loads(base64.b64decode(output_data))
         assert decoded_output_data == expected_output
@@ -68,7 +69,7 @@ def test_read_write_result_file_no_dict(self, tmpdir, file_svc):
         output_encoded = str(b64encode(output.encode()), 'utf-8')
         file_svc.write_result_file(link_id=link_id, output=output_encoded, location=tmpdir)
 
-        expected_output = {'stdout': output, 'stderr': ''}
+        expected_output = {'stdout': output, 'stderr': '', 'exit_code': ''}
         output_data = file_svc.read_result_file(link_id=link_id, location=tmpdir)
         decoded_output_data = json.loads(base64.b64decode(output_data))
         assert decoded_output_data == expected_output
@@ -78,7 +79,7 @@ def test_read_write_result_file_no_base64(self, tmpdir, file_svc):
         output = 'output testing unit'
         file_svc.write_result_file(link_id=link_id, output=output, location=tmpdir)
 
-        expected_output = {'stdout': output, 'stderr': ''}
+        expected_output = {'stdout': output, 'stderr': '', 'exit_code': ''}
         output_data = file_svc.read_result_file(link_id=link_id, location=tmpdir)
         decoded_output_data = json.loads(base64.b64decode(output_data))
         assert decoded_output_data == expected_output