Make it so tests' PCAPs don't automatically get netbox enriched by de…

…fault, it has to be explicitly specified
idaholab · Dec 4, 2024 · 2f7084e · 2f7084e
1 parent 5b5bd14
commit 2f7084e
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 24 deletions.
diff --git a/README.md b/README.md
@@ -155,7 +155,6 @@ $ malcolm-test \
     --rm \
     --image malcolm-testing \
     --vm-provision-os false \
-    --vm-provision-malcolm false \
     --pcap-path /path/to/Malcolm-Test-PCAP
 ====================== test session starts ======================
 platform linux -- Python 3.13.0, pytest-8.3.3, pluggy-1.5.0
@@ -175,7 +174,6 @@ Explanation of arguments:
 * `--rm`: discard the VM after running the tests
 * `--image malcolm-testing`: use the `malcolm-testing` image already built instead of building one from scratch
 * `--vm-provision-os false`: since the image is already provisioned, skip VM provisioning steps
-* `--vm-provision-malcolm false`: since the image already configured Malcolm, skip Malcolm provisioning steps
 * `--pcap-path /path/to/Malcolm-Test-PCAP`: specify the path to the upload artifacts used by the tests
 
 #### <a name="LongRunning"></a> Starting a Malcolm VM and Leaving It Running For Subsequent Test Runs
@@ -190,7 +188,6 @@ $ malcolm-test \
     --rm \
     --image malcolm-testing \
     --vm-provision-os false \
-    --vm-provision-malcolm false \
     --run-tests false
 ```
 
@@ -200,7 +197,6 @@ Explanation of arguments:
 * `--rm`: discard the VM after shutting down
 * `--image malcolm-testing`: use the `malcolm-testing` image already built instead of building one from scratch
 * `--vm-provision-os false`: since the image is already provisioned, skip VM provisioning steps
-* `--vm-provision-malcolm false`: since the image already configured Malcolm, skip Malcolm provisioning steps
 * `--run-tests false`: don't run the tests with this instance of `malcolm-test`, just start Malcolm and wait
 
 In another shell session, connecting to the existing running Malcolm instance for a test run:
@@ -216,7 +212,6 @@ $ malcolm-test \
   --sleep 0 \
   --existing-vm malcolm-nearby-satyr \
   --vm-provision-os false \
-  --vm-provision-malcolm false \
   --run-tests \
   --pcap-path /path/to/Malcolm-Test-PCAP
 ====================== test session starts ======================
@@ -237,7 +232,7 @@ collected 4 items
 * `--sleep 0`: since Malcolm was already started, there is no need to sleep before beginning test execution
 * `--existing-vm malcolm-nearby-satyr`: specify the name of the existing running VM obtained by `virter vm ls`
 * `--vm-provision-os false`: since the image is already provisioned, skip VM provisioning steps
-* `--vm-provision-malcolm false`: since the image already configured Malcolm, skip Malcolm provisioning steps
+* `--vm-provision-malcolm false`: since the Malcolm instance is already configured, skip Malcolm provisioning steps
 * `--run-tests`: run the test suite
 * `--pcap-path /path/to/Malcolm-Test-PCAP`: specify the path to the upload artifacts used by the tests
 
@@ -300,6 +295,8 @@ def test_malcolm_pcap_hash(
 
 As PCAP files are uploaded to Malcolm by `malcolm-test`, they are [hashed](https://docs.python.org/3/library/hashlib.html#hashlib.shake_256) and stored in `pcap_hash_map` which maps the PCAP filename to its hash. The PCAP file is renamed to the hex-representation of the hash digest and the file extension (e.g., if the contents of `foobar.pcap` hashed to `52b92cdcec4af0e1`, the file would be uploaded as `52b92cdcec4af0e1.pcap`). This is done to ensure that conflicting PCAP filenames among different tests are resolved prior to processing. Since Malcolm automatically [assigns tags](https://idaholab.github.io/Malcolm/docs/upload.html#Tagging) to uploaded PCAP files, this hash should be used as a filter for the `tags` field for any network log-based queries for data related to that PCAP file. This way your tests can have reproducible outputs without being affected by PCAPs for other tests.
 
+By default, `malcolm-test` instructs Malcolm to skip [NetBox enrichment](https://idaholab.github.io/Malcolm/docs/upload.html#UploadNetBoxSite) for PCAPs found in `UPLOAD_ARTIFACTS`. To make a test perform NetBox enrichment for its PCAP (using the Malcolm instance's default NetBox site), set `NETBOX_ENRICH` to `True` in the test source.
+
 See the following tests for examples of how to access and use these fixtures:
 
 * [test_malcolm_response.py](src/maltest/tests/test_malcolm_response.py) - querying the [Malcolm API](https://idaholab.github.io/Malcolm/docs/api.html#API) using the [Requests](https://requests.readthedocs.io/en/latest/) library

diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "malcolm-test"
-version = "0.6.3"
+version = "0.6.4"
 authors = [
   { name="Seth Grover", email="[email protected]" },
 ]

diff --git a/src/maltest/maltest.py b/src/maltest/maltest.py
@@ -405,15 +405,11 @@ def main():
                 )
 
                 # for the tests we're about to run, get the set of PCAP files referenced and upload them to Malcolm
-                pcaps = []
+                pcaps = {}
                 if testSetPreExec.collected:
                     pcaps = testSetPreExec.PCAPsReferenced()
-                    logging.debug(
-                        json.dumps(
-                            {'tests': list(testSetPreExec.collected), 'pcaps': list(testSetPreExec.PCAPsReferenced())}
-                        )
-                    )
-                    for pcapFile in pcaps:
+                    logging.debug(json.dumps({'tests': list(testSetPreExec.collected), 'pcaps': pcaps}))
+                    for pcapFile, pcapAttrs in pcaps.items():
                         pcapFilespec = pcapFile if os.path.isabs(pcapFile) else os.path.join(args.pcapPath, pcapFile)
                         pcapFileParts = os.path.splitext(pcapFilespec)
                         if pcapHash := shakey_file_hash(pcapFilespec):
@@ -427,7 +423,7 @@ def main():
                                     copyCode = malcolmVm.CopyFile(
                                         pcapFilespec,
                                         # TODO: Assuming the Malcolm directory like this might not be very robust
-                                        f'/home/{args.vmImageUsername}/Malcolm/pcap/upload/{pcapNewName}',
+                                        f"/home/{args.vmImageUsername}/Malcolm/pcap/upload/{'' if (pcapAttrs['netbox'] == True) else 'NBSITEID0,'}{pcapNewName}",
                                         tolerateFailure=True,
                                     )
                                     if copyCode == 0:

diff --git a/src/maltest/utils.py b/src/maltest/utils.py
@@ -43,6 +43,9 @@
 # tests should define UPLOAD_ARTIFACTS for files (e.g., PCAPs) they need uploaded to Malcolm
 UPLOAD_ARTIFACT_LIST_NAME = 'UPLOAD_ARTIFACTS'
 
+# tests should define NETBOX_ENRICH=True for PCAPs to be netbox-enriched (default false)
+NETBOX_ENRICH_BOOL_NAME = "NETBOX_ENRICH"
+
 # parameters for checking that a Malcolm instance is ready
 MALCOLM_READY_TIMEOUT_SECONDS = 600
 MALCOLM_READY_CHECK_PERIOD_SECONDS = 30
@@ -343,7 +346,7 @@ class MalcolmTestCollection(object):
         testSetPreExec = MalcolmTestCollection()
         pytest.main(['--collect-only', '-p', 'no:terminal'], plugins=[testSetPreExec])
         if testSetPreExec.collected:
-            for pcapFile in testSetPreExec.PCAPsReferenced():
+            for pcapFile, pcapAttrs in testSetPreExec.PCAPsReferenced().items():
                 ...
 
     This allows for determining artifacts used by tests prior to running the tests themselves.
@@ -385,24 +388,45 @@ def PCAPsReferenced(self):
         PCAPsReferenced: Process collected pytest test files by parsing them and looking for UPLOAD_ARTIFACTS,
         which are then collected into a set of all referenced artifacts to be uploaded.
         This can be used to upload these artifacts to Malcolm and make sure they finish processing prior to
-        running the tests that depend on them.
+        running the tests that depend on them. Other variables looked for include:
+            NETBOX_ENRICH - True to indicate that the PCAP should be NetBox-enriched, otherwise it will not be;
+                            sets "netbox" in PCAP sub-dict to True or False
 
         Args:
             None
 
         Returns:
-            A set of all files defined in UPLOAD_ARTIFACTS for all tests to be run
+            A dict of all files defined in UPLOAD_ARTIFACTS for all tests to be run, where the
+            key is the PCAP name and the value is a dict containing other relevant information, e.g.:
+            {
+                "foobar.pcap" : {
+                    "netbox": True
+                },
+                "barbaz.pcap" : {
+                    "netbox": False
+                },
+            }
         """
-        result = list()
+        result = dict()
         for testPyPath in self.collected:
             try:
+                testArtifactList = list()
+                testNetBoxEnrich = False
                 with open(testPyPath, "r") as f:
                     testPyContent = f.read()
                 for node in ast.walk(ast.parse(testPyContent)):
                     if isinstance(node, ast.Assign):
                         for target in node.targets:
-                            if isinstance(target, ast.Name) and (target.id == UPLOAD_ARTIFACT_LIST_NAME):
-                                result.append(ast.literal_eval(node.value))
+                            if isinstance(target, ast.Name):
+                                if target.id == UPLOAD_ARTIFACT_LIST_NAME:
+                                    testArtifactList.append(ast.literal_eval(node.value))
+                                elif target.id == NETBOX_ENRICH_BOOL_NAME:
+                                    testNetBoxEnrich = mmguero.str2bool(str(ast.literal_eval(node.value)))
+                for artifact in list(set(mmguero.Flatten(testArtifactList))):
+                    if artifact not in result:
+                        result[artifact] = defaultdict(lambda: None)
+                    if testNetBoxEnrich:
+                        result[artifact]["netbox"] = True
             except FileNotFoundError:
                 self.logger.error(f"Error: '{testPyPath}' not found")
             except SyntaxError:
@@ -411,7 +435,7 @@ def PCAPsReferenced(self):
                 self.logger.error(f"Error: Unable to evaulate '{variable_name}' in '{testPyPath}': {ve}")
             except Exception as e:
                 self.logger.error(f"Error: '{testPyPath}': {e}")
-        return set(mmguero.Flatten(result))
+        return result
 
 
 ###################################################################################################
@@ -779,6 +803,8 @@ def ArkimeAlreadyHasFile(
                             **self.dbObjs.DatabaseInitArgs,
                         ),
                         index=ARKIME_FILES_INDEX,
+                        # note that the query here is *foobar.pcap, so it will match both
+                        #   foobar.pcap and NBSITEID0,foobar.pcap
                     ).query("wildcard", name=f"*{os.path.basename(filename)}")
                     response = s.execute()
                     for hit in response:

diff --git a/src/maltest/virter/malcolm-init/01-clone-install.toml b/src/maltest/virter/malcolm-init/01-clone-install.toml
@@ -76,8 +76,8 @@ VIRUSTOTAL_API_KEY = ""
 FILE_SCAN_RULE_UPDATE = "false"
 NETBOX = "true"
 NETBOX_ENRICH = "true"
-NETBOX_AUTOPOPULATE = "false"
-NETBOX_AUTO_PREFIXES = "false"
+NETBOX_AUTOPOPULATE = "true"
+NETBOX_AUTO_PREFIXES = "true"
 NETBOX_SITE_NAME = ""
 LIVE_CAPTURE_IFACE = ""
 LIVE_CAPTURE_FILTER = ""