Limit tab completion of scan command to scans with installed componen…

…ts (#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.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 @@ def __init__(self, *args, **kwargs):
         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,11 +1,23 @@
 import sys
+import pickle
+import typing
 import inspect
 import pkgutil
 import importlib
 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"):
-                    # now we only care about classes that end in [Ss]can
-                    scans[subname].append(f"{__package__}.{name}")
+            for sub_name, sub_obj in inspect.getmembers(obj):
+                # now we only care about classes that end in [Ss]can
+                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 ..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 @@ def __init__(self, *args, **kwargs):
         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 .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 @@ def __init__(self, *args, **kwargs):
         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 ...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 @@ def __init__(self, *args, **kwargs):
         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.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 @@ def __init__(self, *args, **kwargs):
         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 @@
 
 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 @@ def __init__(self, *args, **kwargs):
         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 @@ def __init__(self, *args, **kwargs):
         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 .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 @@ def __init__(self, *args, **kwargs):
         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 @@
 
 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 @@ def __init__(self, *args, **kwargs):
         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 @@
 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())
-
-    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",
-    ]
+    scans = get_scans()
 
-    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 @@ def setup_method(self):
         )
         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"