Added initial callback payload support

Author: SteveMcGrath

.gitignore

@@ -7,3 +7,4 @@ __pycache__
 .vars
 .secrets
 requirements.txt
+/example

tenint/__init__.py

@@ -2,4 +2,4 @@
 from .models.configuration import Configuration, Settings  # noqa F401
 from .connector import Connector  # noqa F401
 
-__version__ = '0.9.1'
+__version__ = '0.9.2'

tenint/connector.py

@@ -20,9 +20,12 @@
 from typer import Exit, Option, Typer
 from typer.main import get_command_name
 
+from .models.callback import CallbackResponse
 from .models.configuration import Configuration, Settings
 from .models.credentials import Credential
 
+logger = logging.getLogger('tenint.connector')
+
 
 class ConfigurationError(Exception):
     pass
@@ -61,7 +64,6 @@ def __init__(
     ):
         self.app = Typer(add_completion=False)
         self.console = Console()
-
         self.settings = settings
         self.defaults = (
             defaults if defaults else settings.model_construct().model_dump()
@@ -215,25 +217,53 @@ def cmd_run(
         """
         logging.basicConfig(
             level=log_level.upper(),
-            format='%(asctime)s %(levelname)s %(message)s',
+            format='%(name)s: %(message)s',
             datefmt='%Y-%m-%d %H:%M:%S',
-            handlers=[RichHandler(console=self.console, rich_tracebacks=True)],
+            handlers=[
+                RichHandler(
+                    console=self.console,
+                    rich_tracebacks=True,
+                    omit_repeated_times=False,
+                )
+            ],
         )
         status_code = 0
+        resp = None
 
+        # Attempt to run the connector job and handle any errors that may be thrown
+        # in a graceful way.
         try:
             config = self.fetch_config(json_data, filename)
-            self.main(config=config, since=since)
-        except (ValidationError, ConfigurationError) as exc:
-            self.console.print(exc)
+            resp = self.main(config=config, since=since)
+        except (ValidationError, ConfigurationError) as _:
+            logging.error('Invalid configuration presented', exc_info=True)
             status_code = 2
         except Exception as _:
-            self.console.print_exception()
+            logging.error('Job run failed with an error', exc_info=True)
             status_code = 1
 
+        # Set the Callback payload to the job response if the response is a dictionary
+        try:
+            payload = CallbackResponse(exit_code=status_code, **resp).model_dump(
+                mode='json'
+            )
+        except (ValidationError, TypeError) as _:
+            logger.error(f'Job response format is invalid: {resp=}')
+            payload = CallbackResponse(exit_code=status_code).model_dump(mode='json')
+
+        # If a callback and a job id have been set, then we will initiate a callback
+        # to the job manager with the response payload of the job to inform the manager
+        # that we have finished.
         if job_id and callback_url:
-            requests.post(callback_url, headers={'X-Job-ID': job_id})
-            self.console.log(f'Called back to {callback_url=} with {job_id=}')
+            requests.post(
+                callback_url,
+                headers={'X-Job-ID': job_id},
+                json=payload,
+            )
+            logger.info(f'Called back to {callback_url=} with {job_id=} and {payload=}')
         else:
-            self.console.print('Did not initiate a callback!', style='bold orange1')
+            logger.warning('Did not initiate a callback!')
+        self.console.print(f'callback={payload}')
+
+        # Exit the connector with the status code from the runner.
         raise Exit(code=status_code)

tenint/models/callback.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel
+
+
+class CallbackDataTypeCounts(BaseModel):
+    sent: int = 0
+    modified: int = 0
+    failed: int = 0
+
+
+class CallbackCounters(BaseModel):
+    assets: CallbackDataTypeCounts = CallbackDataTypeCounts()
+    findings: CallbackDataTypeCounts = CallbackDataTypeCounts()
+
+
+class CallbackResponse(BaseModel):
+    exit_code: int
+    counts: CallbackCounters = CallbackCounters()

tenint/templates/connector.py

@@ -22,12 +22,13 @@ class AppSettings(Settings):
 
 
 @connector.job
-def main(config: AppSettings, since: int | None = None):
+def main(config: AppSettings, since: int | None = None) -> dict:
     log.debug('This is a debug test')
     log.info('this is an info test')
     log.warning('this is a warning')
     log.error('this is an error')
     print('hello world')
+    return {'counts': {'assets': {'sent': 0}}}
 
 
 if __name__ == '__main__':

tests/test_connector.py

@@ -99,13 +99,19 @@ def test_connector_validation(AppSettings):
     assert 'Not yet implemented' in result.stdout
 
 
-def test_connector_run(AppSettings, main):
+def test_connector_run(AppSettings, main, caplog):
     connector = Connector(settings=AppSettings)
     connector.job(main)
-    result = runner.invoke(connector.app, ['run', '-j', '{"is_bool": true}'])
-    assert result.exit_code == 0
-    assert 'hello world' in result.stdout
-    assert 'Did not initiate a callback!' in result.stdout
+    with caplog.at_level(logging.INFO):
+        result = runner.invoke(connector.app, ['run', '-j', '{"is_bool": true}'])
+        assert result.exit_code == 0
+        assert 'hello world' in result.stdout
+        assert 'This is a debug test' not in caplog.text
+        assert 'This is an info test' in caplog.text
+        assert 'This is an error test' in caplog.text
+        assert 'Job response format is invalid' in caplog.text
+        assert 'Did not initiate a callback!' in caplog.text
+        assert 'callback={' in result.stdout
 
 
 def test_connector_run_config_fail(AppSettings, main):
@@ -115,39 +121,42 @@ def test_connector_run_config_fail(AppSettings, main):
     assert result.exit_code == 2
 
 
-def test_connector_run_fail(AppSettings):
+def test_connector_run_fail(AppSettings, caplog):
     connector = Connector(settings=AppSettings)
 
     @connector.job
     def failmain(config: dict, since: None = None):
         raise Exception('I have failed')
 
-    result = runner.invoke(connector.app, ['run', '-j', '{"is_bool": true}'])
-    assert result.exit_code == 1
-    assert 'I have failed' in result.stdout
+    with caplog.at_level(logging.INFO):
+        result = runner.invoke(connector.app, ['run', '-j', '{"is_bool": true}'])
+        assert result.exit_code == 1
+        assert 'I have failed' in caplog.text
 
 
 @responses.activate
-def test_connector_callback(AppSettings, main):
+def test_connector_callback(AppSettings, main, caplog):
     responses.post(
         'http://callback-url.local/callback',
         match=[matchers.header_matcher({'X-Job-ID': 'abcdef'})],
     )
     connector = Connector(settings=AppSettings)
     connector.job(main)
-    result = runner.invoke(
-        connector.app,
-        [
-            'run',
-            '-j',
-            '{"is_bool": true}',
-            '-J',
-            'abcdef',
-            '-c',
-            'http://callback-url.local/callback',
-        ],
-    )
-    assert result.exit_code == 0
-    assert 'Called back' in result.stdout
-    assert "callback_url='http://callback-url.local/callback'" in result.stdout
-    assert "job_id='abcdef'" in result.stdout
+
+    with caplog.at_level(logging.INFO):
+        result = runner.invoke(
+            connector.app,
+            [
+                'run',
+                '-j',
+                '{"is_bool": true}',
+                '-J',
+                'abcdef',
+                '-c',
+                'http://callback-url.local/callback',
+            ],
+        )
+        assert result.exit_code == 0
+        assert 'Called back' in caplog.text
+        assert "callback_url='http://callback-url.local/callback'" in caplog.text
+        assert "job_id='abcdef'" in caplog.text