Limit tab completion of scan command to scans with installed components (#74)

* scans without installed requirements dont tab-complete

* added meets_requirements functions to classes

* updated tests

* added check for None case

pipeline/recon/amass.py

@@ -7,8 +7,9 @@ from luigi.util import inherits
 from luigi.contrib.sqla import SQLAlchemyTarget
 
 import pipeline.models.db_manager
-from .targets import TargetList
 from ..tools import tools
+from .targets import TargetList
+from .helpers import get_tool_state
 from ..models.target_model import Target
 
 
@@ -48,6 +49,15 @@ class AmassScan(luigi.Task):
         self.db_mgr = pipeline.models.db_manager.DBManager(db_location=self.db_location)
         self.results_subfolder = (Path(self.results_dir) / "amass-results").expanduser().resolve()
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["amass"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ AmassScan depends on TargetList to run.
 

pipeline/recon/helpers.py

@@ -1,4 +1,6 @@
 import sys
+import pickle
+import typing
 import inspect
 import pkgutil
 import importlib
@@ -6,6 +8,16 @@ import ipaddress
 from pathlib import Path
 from collections import defaultdict
 
+from ..recon.config import defaults
+
+
+def get_tool_state() -> typing.Union[dict, None]:
+    """ Load current tool state from disk. """
+    tools = Path(defaults.get("tools-dir")) / ".tool-dict.pkl"
+
+    if tools.exists():
+        return pickle.loads(tools.read_bytes())
+
 
 def get_scans():
     """ Iterates over the recon package and its modules to find all of the classes that end in [Ss]can.
@@ -42,10 +54,19 @@ def get_scans():
         if inspect.ismodule(obj) and not name.startswith("_"):
             # we're only interested in modules that don't begin with _ i.e. magic methods __len__ etc...
 
-            for subname, subobj in inspect.getmembers(obj):
-                if inspect.isclass(subobj) and subname.lower().endswith("scan"):
+            for sub_name, sub_obj in inspect.getmembers(obj):
                 # now we only care about classes that end in [Ss]can
-                    scans[subname].append(f"{__package__}.{name}")
+                if inspect.isclass(sub_obj) and sub_name.lower().endswith("scan"):
+                    # final check, this ensures that the tools necessary to AT LEAST run this scan are present
+                    # does not consider upstream dependencies
+                    try:
+                        if not sub_obj.meets_requirements():
+                            continue
+                    except AttributeError:
+                        # some scan's haven't implemented meets_requirements yet, silently allow them through
+                        pass
+
+                    scans[sub_name].append(f"{__package__}.{name}")
 
     return scans
 

pipeline/recon/masscan.py

@@ -14,6 +14,7 @@ from .amass import ParseAmassOutput
 from ..models.port_model import Port
 from ..models.ip_address_model import IPAddress
 
+from .helpers import get_tool_state
 from .config import top_tcp_ports, top_udp_ports, defaults, web_ports
 
 
@@ -63,6 +64,15 @@ class MasscanScan(luigi.Task):
         self.db_mgr = pipeline.models.db_manager.DBManager(db_location=self.db_location)
         self.results_subfolder = (Path(self.results_dir) / "masscan-results").expanduser().resolve()
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["masscan"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def output(self):
         """ Returns the target output for this task.
 

pipeline/recon/nmap.py

@@ -16,6 +16,7 @@ from .config import defaults
 from .helpers import get_ip_address_version, is_ip_address
 
 from ..tools import tools
+from .helpers import get_tool_state
 from ..models.port_model import Port
 from ..models.nse_model import NSEResult
 from ..models.target_model import Target
@@ -241,6 +242,15 @@ class SearchsploitScan(luigi.Task):
         super().__init__(*args, **kwargs)
         self.db_mgr = pipeline.models.db_manager.DBManager(db_location=self.db_location)
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["searchsploit"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ Searchsploit depends on ThreadedNmap to run.
 

pipeline/recon/web/aquatone.py

@@ -13,6 +13,7 @@ from ..config import defaults
 from ...tools import tools
 
 import pipeline.models.db_manager
+from ..helpers import get_tool_state
 from ...models.port_model import Port
 from ...models.header_model import Header
 from ...models.endpoint_model import Endpoint
@@ -63,6 +64,15 @@ class AquatoneScan(luigi.Task):
         self.db_mgr = pipeline.models.db_manager.DBManager(db_location=self.db_location)
         self.results_subfolder = Path(self.results_dir) / "aquatone-results"
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["aquatone"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ AquatoneScan depends on GatherWebTargets to run.
 

pipeline/recon/web/gobuster.py

@@ -10,9 +10,10 @@ from luigi.util import inherits
 from luigi.contrib.sqla import SQLAlchemyTarget
 
 import pipeline.models.db_manager
-from .targets import GatherWebTargets
-from ..config import defaults
 from ...tools import tools
+from ..config import defaults
+from ..helpers import get_tool_state
+from .targets import GatherWebTargets
 from ...models.endpoint_model import Endpoint
 from ..helpers import get_ip_address_version, is_ip_address
 
@@ -64,6 +65,15 @@ class GobusterScan(luigi.Task):
         self.db_mgr = pipeline.models.db_manager.DBManager(db_location=self.db_location)
         self.results_subfolder = Path(self.results_dir) / "gobuster-results"
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["recursive-gobuster", "gobuster"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ GobusterScan depends on GatherWebTargets to run.
 

pipeline/recon/web/subdomain_takeover.py

@@ -9,8 +9,9 @@ from luigi.contrib.sqla import SQLAlchemyTarget
 
 import pipeline.models.db_manager
 from ...tools import tools
-from .targets import GatherWebTargets
 from ..config import defaults
+from ..helpers import get_tool_state
+from .targets import GatherWebTargets
 
 
 @inherits(GatherWebTargets)
@@ -52,6 +53,15 @@ class TKOSubsScan(luigi.Task):
         self.results_subfolder = (Path(self.results_dir) / "tkosubs-results").expanduser().resolve()
         self.output_file = self.results_subfolder / "tkosubs.csv"
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["tko-subs"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ TKOSubsScan depends on GatherWebTargets to run.
 
@@ -174,6 +184,15 @@ class SubjackScan(luigi.Task):
         self.results_subfolder = (Path(self.results_dir) / "subjack-results").expanduser().resolve()
         self.output_file = self.results_subfolder / "subjack.txt"
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["subjack"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ SubjackScan depends on GatherWebTargets to run.
 

pipeline/recon/web/waybackurls.py

@@ -8,6 +8,7 @@ from luigi.contrib.sqla import SQLAlchemyTarget
 
 from .targets import GatherWebTargets
 from ...tools import tools
+from ..helpers import get_tool_state
 from ...models.endpoint_model import Endpoint
 
 import pipeline.models.db_manager
@@ -48,6 +49,15 @@ class WaybackurlsScan(luigi.Task):
         self.db_mgr = pipeline.models.db_manager.DBManager(db_location=self.db_location)
         self.results_subfolder = Path(self.results_dir) / "waybackurls-results"
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["waybackurls"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ WaybackurlsScan depends on GatherWebTargets to run.
 

pipeline/recon/web/webanalyze.py

@@ -12,8 +12,9 @@ from luigi.contrib.sqla import SQLAlchemyTarget
 
 import pipeline.models.db_manager
 from ...tools import tools
-from .targets import GatherWebTargets
 from ..config import defaults
+from ..helpers import get_tool_state
+from .targets import GatherWebTargets
 from ...models.technology_model import Technology
 from ..helpers import get_ip_address_version, is_ip_address
 
@@ -59,6 +60,15 @@ class WebanalyzeScan(luigi.Task):
         self.db_mgr = pipeline.models.db_manager.DBManager(db_location=self.db_location)
         self.results_subfolder = Path(self.results_dir) / "webanalyze-results"
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["webanalyze"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ WebanalyzeScan depends on GatherWebTargets to run.
 

pipeline/recon/wrappers.py

@@ -2,6 +2,7 @@ import luigi
 from luigi.util import inherits
 
 from .nmap import SearchsploitScan
+from .helpers import get_tool_state
 from .web import AquatoneScan, GobusterScan, SubjackScan, TKOSubsScan, WaybackurlsScan, WebanalyzeScan
 
 
@@ -27,6 +28,26 @@ class FullScan(luigi.WrapperTask):
         results_dir: specifes the directory on disk to which all Task results are written
     """
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = [
+            "amass",
+            "aquatone",
+            "masscan",
+            "tko-subs",
+            "recursive-gobuster",
+            "searchsploit",
+            "subjack",
+            "gobuster",
+            "webanalyze",
+            "waybackurls",
+        ]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ FullScan is a wrapper, as such it requires any Tasks that it wraps. """
         args = {
@@ -90,6 +111,15 @@ class HTBScan(luigi.WrapperTask):
         results_dir: specifes the directory on disk to which all Task results are written
     """
 
+    @staticmethod
+    def meets_requirements():
+        """ Reports whether or not this scan's needed tool(s) are installed or not """
+        needs = ["aquatone", "masscan", "recursive-gobuster", "searchsploit", "gobuster", "webanalyze"]
+        tools = get_tool_state()
+
+        if tools:
+            return all([tools.get(x).get("installed") is True for x in needs])
+
     def requires(self):
         """ HTBScan is a wrapper, as such it requires any Tasks that it wraps. """
         args = {

tests/test_recon/test_helpers.py

@@ -1,48 +1,34 @@
 import pytest
 
 from pipeline.recon.helpers import get_ip_address_version, get_scans, is_ip_address
+from pipeline.recon import AmassScan, MasscanScan, FullScan, HTBScan, SearchsploitScan, ThreadedNmapScan
+from pipeline.recon.web import GobusterScan, SubjackScan, TKOSubsScan, AquatoneScan, WaybackurlsScan, WebanalyzeScan
 
 
 def test_get_scans():
-    scans = get_scans()
 
     scan_names = [
-        "AmassScan",
-        "GobusterScan",
-        "MasscanScan",
-        "SubjackScan",
-        "TKOSubsScan",
-        "AquatoneScan",
-        "FullScan",
-        "HTBScan",
-        "SearchsploitScan",
-        "ThreadedNmapScan",
-        "WebanalyzeScan",
-        "WaybackurlsScan",
+        AmassScan,
+        GobusterScan,
+        MasscanScan,
+        SubjackScan,
+        TKOSubsScan,
+        AquatoneScan,
+        FullScan,
+        HTBScan,
+        SearchsploitScan,
+        ThreadedNmapScan,
+        WebanalyzeScan,
+        WaybackurlsScan,
     ]
 
-    assert len(scan_names) == len(scans.keys())
+    scans = get_scans()
 
-    for name in scan_names:
-        assert name in scans.keys()
-
-    modules = [
-        "pipeline.recon.amass",
-        "pipeline.recon.masscan",
-        "pipeline.recon.nmap",
-        "pipeline.recon.nmap",
-        "pipeline.recon.web",
-        "pipeline.recon.web",
-        "pipeline.recon.web",
-        "pipeline.recon.web",
-        "pipeline.recon.web",
-        "pipeline.recon.web",
-        "pipeline.recon.wrappers",
-        "pipeline.recon.wrappers",
-    ]
-
-    for module in scans.values():
-        assert module[0] in modules
+    for scan in scan_names:
+        if hasattr(scan, "meets_requirements") and scan.meets_requirements():
+            assert scan.__name__ in scans.keys()
+        else:
+            assert scan not in scans.keys()
 
 
 @pytest.mark.parametrize(

tests/test_recon/test_masscan.py

@@ -40,6 +40,9 @@ class TestParseMasscanOutput:
         )
         self.scan.input = lambda: luigi.LocalTarget(masscan_results)
 
+    def teardown_method(self):
+        shutil.rmtree(self.tmp_path)
+
     def test_scan_creates_results_dir(self):
         assert self.scan.results_subfolder == self.tmp_path / "masscan-results"
 

tests/test_shell/test_recon_pipeline_shell.py

@@ -428,12 +428,17 @@ class TestReconShell:
             "webbrowser.open", autospec=True
         ) as mocked_web, patch("selectors.DefaultSelector.register", autospec=True) as mocked_selector, patch(
             "cmd2.Cmd.select"
-        ) as mocked_select:
+        ) as mocked_select, patch(
+            "pipeline.recon-pipeline.get_scans"
+        ) as mocked_scans:
+
             mocked_select.return_value = "Resume"
             mocked_popen.return_value = process_mock
 
             test_input += f" --results-dir {tmp_path / 'mostuff'}"
 
+            mocked_scans.return_value = {"FullScan": ["pipeline.recon.wrappers"]}
+
             if db_mgr is None:
                 self.shell.do_scan(test_input)
                 assert expected in capsys.readouterr().out

tests/test_web/test_webanalyze.py

@@ -4,8 +4,8 @@ import tempfile
 from pathlib import Path
 from unittest.mock import patch, MagicMock
 
-from pipeline.tools import tools
 from pipeline.recon.web import WebanalyzeScan, GatherWebTargets
+from pipeline.tools import tools
 
 webanalyze_results = Path(__file__).parent.parent / "data" / "recon-results" / "webanalyze-results"