diff --git a/pipeline/recon/amass.py b/pipeline/recon/amass.py index 77b0f8f..04382a5 100644 --- a/pipeline/recon/amass.py +++ b/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. diff --git a/pipeline/recon/helpers.py b/pipeline/recon/helpers.py index bd3007a..bc646fd 100644 --- a/pipeline/recon/helpers.py +++ b/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"): - # 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 diff --git a/pipeline/recon/masscan.py b/pipeline/recon/masscan.py index ad9abb8..2b95f98 100644 --- a/pipeline/recon/masscan.py +++ b/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. diff --git a/pipeline/recon/nmap.py b/pipeline/recon/nmap.py index 2169390..51bacca 100644 --- a/pipeline/recon/nmap.py +++ b/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. diff --git a/pipeline/recon/web/aquatone.py b/pipeline/recon/web/aquatone.py index 4464a82..cdb46c4 100644 --- a/pipeline/recon/web/aquatone.py +++ b/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. diff --git a/pipeline/recon/web/gobuster.py b/pipeline/recon/web/gobuster.py index f906e9e..ee0f003 100644 --- a/pipeline/recon/web/gobuster.py +++ b/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. diff --git a/pipeline/recon/web/subdomain_takeover.py b/pipeline/recon/web/subdomain_takeover.py index aa34dab..57b64bc 100644 --- a/pipeline/recon/web/subdomain_takeover.py +++ b/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. diff --git a/pipeline/recon/web/waybackurls.py b/pipeline/recon/web/waybackurls.py index f09db7c..6d97cdf 100644 --- a/pipeline/recon/web/waybackurls.py +++ b/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. diff --git a/pipeline/recon/web/webanalyze.py b/pipeline/recon/web/webanalyze.py index 15608dd..07b2f35 100644 --- a/pipeline/recon/web/webanalyze.py +++ b/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. diff --git a/pipeline/recon/wrappers.py b/pipeline/recon/wrappers.py index 327fbd3..2ee1ff3 100644 --- a/pipeline/recon/wrappers.py +++ b/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 = { diff --git a/tests/test_recon/test_helpers.py b/tests/test_recon/test_helpers.py index 5e7aee4..97f2c62 100644 --- a/tests/test_recon/test_helpers.py +++ b/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( diff --git a/tests/test_recon/test_masscan.py b/tests/test_recon/test_masscan.py index dca93ab..dfa8286 100644 --- a/tests/test_recon/test_masscan.py +++ b/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" diff --git a/tests/test_shell/test_recon_pipeline_shell.py b/tests/test_shell/test_recon_pipeline_shell.py index d5bd071..f86e51e 100644 --- a/tests/test_shell/test_recon_pipeline_shell.py +++ b/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 diff --git a/tests/test_web/test_webanalyze.py b/tests/test_web/test_webanalyze.py index 2861e54..67803c1 100644 --- a/tests/test_web/test_webanalyze.py +++ b/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"