Merge pull request #50 from tristanlatr/dev3

Change daemon mode
tristanlatr · Feb 16, 2021 · 2461888 · 2461888
2 parents e92bfb3 + 1cb6dbf
commit 2461888
Show file tree

Hide file tree

Showing 9 changed files with 187 additions and 110 deletions.
diff --git a/setup.py b/setup.py
@@ -1,10 +1,7 @@
 #! /usr/bin/env python3
 from setuptools import setup, find_packages
-import sys
-if sys.version_info[0] < 3: 
-    raise RuntimeError("You must use Python 3")
-# The directory containing this file
 import pathlib
+# The directory containing this file
 HERE = pathlib.Path(__file__).parent
 # About the project
 ABOUT = {}
@@ -19,13 +16,23 @@
     version             =   ABOUT['__version__'],
     packages            =   find_packages(exclude=('tests')), 
     entry_points        =   {'console_scripts': ['wpwatcher = wpwatcher.cli:main'],},
-    classifiers         =   ["Programming Language :: Python :: 3"],
+    classifiers         =   [
+        "Development Status :: 5 - Production/Stable",
+        "Intended Audience :: Information Technology",
+        "Environment :: Console",
+        "Topic :: Security",
+        "Topic :: Utilities",
+        "Topic :: System :: Monitoring",
+        "Programming Language :: Python :: 3",
+        "Typing :: Typed",
+        "License :: OSI Approved :: Apache Software License", ],
     license             =   ABOUT['__license__'],
     long_description    =   README,
     long_description_content_type   =   "text/markdown",
-    install_requires    =   ['wpscan-out-parse>=1.8.1', 'filelock'],
+    python_requires     =   '>=3.6',
+    install_requires    =   ['wpscan-out-parse>=1.8.1', 'filelock', ],
     extras_require      =   {'syslog' : ['rfc5424-logging-handler', 'cefevent'],
-                             'docs': ["Sphinx", "sphinx_rtd_theme", "recommonmark"], 
+                             'docs': ["Sphinx", "recommonmark"], 
                              'dev': ["pytest", "pytest-cov", "codecov", "coverage", "tox", "mypy"]},
     keywords            =   ABOUT['__keywords__'],
 )
diff --git a/tests/core_test.py b/tests/core_test.py
@@ -1,13 +1,14 @@
 import unittest
-import os
 import shlex
+from datetime import  timedelta
 from . import DEFAULT_CONFIG
 from wpwatcher.core import WPWatcher
 from wpwatcher.config import Config
 from wpwatcher.scan import Scanner
 from wpwatcher.email import EmailSender
 from wpwatcher.wpscan import WPScanWrapper
-
+from wpwatcher.daemon import Daemon
+from wpwatcher.utils import timeout
 class T(unittest.TestCase):
 
     def test_interrupt(self):
@@ -31,11 +32,18 @@ def test_asynch_exec(self):
 
     def test_daemon(self):
         # test daemon_loop_sleep and daemon mode
-        pass
+
+        conf = Config.fromstring(DEFAULT_CONFIG)
+        conf['asynch_workers']+=1
+        daemon = Daemon(conf)
+
+        daemon.loop(ttl=timedelta(seconds=5))
+
+        self.assertTrue(not any([r.status() != 'ERROR' for r in daemon.wpwatcher.new_reports]))
+        self.assertGreater(len(daemon.wpwatcher.new_reports), 1)
+
 
     def test_fail_fast(self):
         pass
 
-    def test_run_scans_and_notify(self):
-        # test returned results
-        pass
+
diff --git a/wpwatcher/cli.py b/wpwatcher/cli.py
@@ -61,13 +61,15 @@ def main(_args: Optional[Sequence[Text]] = None) -> None:
 
     # If daemon lopping
     if configuration["daemon"]:
+
         # Run 4 ever
-        Daemon(configuration)
+        daemon = Daemon(configuration)
+        daemon.loop()
+
     else:
         # Run scans and quit
-        # Create main object
         wpwatcher = WPWatcher(configuration)
-        exit_code, _ = wpwatcher.run_scans()
+        exit_code, reports = wpwatcher.run_scans()
         exit(exit_code)
 
 
@@ -107,8 +109,8 @@ def show(urlpart: str, filepath: Optional[str] = None, daemon: bool = False) ->
 
 def version() -> None:
     """Print version and contributors"""
-    log.info(f"Version:\t\t{__version__}")
-    log.info(f"Authors:\t\t{__author__}")
+    print(f"Version:\t\t{__version__}")
+    print(f"Authors:\t\t{__author__}")
     exit(0)
 
 

diff --git a/wpwatcher/config.py b/wpwatcher/config.py
@@ -431,7 +431,7 @@ def find_files(
         default_content: str = "",
         create: bool = False,
     ) -> List[str]:
-        """Find existent files based on folders name and file names.
+        """Find existent files or folders based on folders name and file names.
 
         Arguments:
         - `env_location`: list of environment variable to use as a base path. Exemple: ['HOME', 'XDG_CONFIG_HOME', 'APPDATA', 'PWD']
@@ -451,9 +451,9 @@ def find_files(
                     potential_paths.append(os.path.join(os.environ[env_var], file_path))
         if not env_loc_exists:
             raise RuntimeError(f"Cannot find any of the env locations {env_location}. ")
-        # If file exist, add to list
+        # If file or folder exist, add to list
         for p in potential_paths:
-            if os.path.isfile(p):
+            if os.path.isfile(p) or os.path.isdir(p):
                 existent_files.append(p)
         # If no file foud and create=True, init new template config
         if len(existent_files) == 0 and create:

diff --git a/wpwatcher/core.py b/wpwatcher/core.py
@@ -30,16 +30,16 @@ class WPWatcher:
 
     .. python::
 
-        from wpwatcher.config import WPWatcherConfig
+        from wpwatcher.config import Config
         from wpwatcher.core import WPWatcher
-        config = WPWatcherConfig.fromenv()
+        config = Config.fromenv()
         config.update({ 'send_infos':   True,
                         'wp_sites':     [   {'url':'exemple1.com'},
                                             {'url':'exemple2.com'}  ],
                         'wpscan_args': ['--format', 'json', '--stealthy']
                     })
         watcher = WPWatcher(config)
-        exit_code, reports = watcher.run_scans_and_notify()
+        exit_code, reports = watcher.run_scans()
         for r in reports:
             print("%s\t\t%s"%( r['site'], r['status'] ))
     """
@@ -64,28 +64,33 @@ def __init__(self, conf: Config):
         # Init scanner
         self.scanner: Scanner = Scanner(conf)
 
-        # Dump config
-        log.info(f"Configuration:{repr(conf)}")
-
         # Save sites
-        self.wp_sites: List[Site] = [
+        conf["wp_sites"] = [
             Site(site_conf) for site_conf in conf["wp_sites"]
         ]
+        self.wp_sites: List[Site] = conf["wp_sites"]
+
 
         # Asynchronous executor
-        self.executor: concurrent.futures.ThreadPoolExecutor = (
+        self._executor: concurrent.futures.ThreadPoolExecutor = (
             concurrent.futures.ThreadPoolExecutor(max_workers=conf["asynch_workers"])
         )
 
         # List of conccurent futures
-        self.futures: List[concurrent.futures.Future] = []  # type: ignore [type-arg]
+        self._futures: List[concurrent.futures.Future] = []  # type: ignore [type-arg]
 
         # Register the signals to be caught ^C , SIGTERM (kill) , service restart , will trigger interrupt()
         signal.signal(signal.SIGINT, self.interrupt)
         signal.signal(signal.SIGTERM, self.interrupt)
 
-        # new reports
-        self.new_reports: ReportCollection
+        self.new_reports = ReportCollection()
+        "New reports, cleared and filled when running `run_scans`."
+
+        self.all_reports = ReportCollection()
+        "All reports an instance of `WPWatcher` have generated using `run_scans`."
+
+        # Dump config
+        log.debug(f"Configuration:{repr(conf)}")
 
     @staticmethod
     def _delete_tmp_wpscan_files() -> None:
@@ -102,64 +107,69 @@ def _delete_tmp_wpscan_files() -> None:
 
     def _cancel_pending_futures(self) -> None:
         """Cancel all asynchronous jobs"""
-        for f in self.futures:
+        for f in self._futures:
             if not f.done():
                 f.cancel()
 
+    def interrupt_scans(self) -> None:
+        """
+        Interrupt the scans and append finished scan reports to self.new_reports
+        """
+        # Cancel all scans
+        self._cancel_pending_futures()  # future scans
+        self.scanner.interrupt()  # running scans
+        self._rebuild_rew_reports()
+
+    def _rebuild_rew_reports(self) -> None:
+        "Recover reports from futures results"
+        self.new_reports = ReportCollection()
+        for f in self._futures:
+            if f.done():
+                try:
+                    self.new_reports.append(f.result())
+                except Exception:
+                    pass
+
     def interrupt(self, sig=None, frame=None) -> None:  # type: ignore [no-untyped-def]
-        """Interrupt sequence"""
+        """Interrupt the program and exit. """
         log.error("Interrupting...")
         # If called inside ThreadPoolExecutor, raise Exeception
         if not isinstance(threading.current_thread(), threading._MainThread):  # type: ignore [attr-defined]
             raise InterruptedError()
 
-        # Cancel all scans
-        self._cancel_pending_futures()  # future scans
-        self.scanner.interrupt()  # running scans
+        self.interrupt_scans()
 
         # Give a 5 seconds timeout to buggy WPScan jobs to finish or ignore them
         try:
-            timeout(5, self.executor.shutdown, kwargs=dict(wait=True))
+            timeout(5, self._executor.shutdown, kwargs=dict(wait=True))
         except TimeoutError:
             pass
 
-        # Recover reports from futures results
-        self.new_reports = ReportCollection()
-        for f in self.futures:
-            if f.done():
-                try:
-                    self.new_reports.append(f.result())
-                except Exception:
-                    pass
-
-        # Display results and quit
-        self._print_new_reports_results()
+        # Display results 
+        log.info(repr(self.new_reports))
         log.info("Scans interrupted.")
+
+        # and quit
         sys.exit(-1)
 
-    def _print_new_reports_results(self) -> None:
-        """Print the result summary for the scanned sites"""
-        new_reports = ReportCollection(n for n in self.new_reports if n)
-        if len(new_reports) > 0:
-            log.info(repr(new_reports))
+
+    def _log_db_reports_infos(self) -> None:
+        if len(self.new_reports) > 0 and repr(self.new_reports) != "No scan report to show":
             if self.wp_reports.filepath != "null":
-                log.info(
-                    f"Updated {len(new_reports)} reports in database: {self.wp_reports.filepath}"
-                )
+                log.info(f"Updated reports in database: {self.wp_reports.filepath}")
             else:
                 log.info("Local database disabled, no reports updated.")
-        else:
-            log.info("No reports updated.")
 
-    def scan_site(self, wp_site: Site) -> Optional[ScanReport]:
+
+    def _scan_site(self, wp_site: Site) -> Optional[ScanReport]:
         """
         Helper method to wrap the scanning process of `WPWatcherScanner.scan_site` and add the following:
+        
         - Find the last report in the database and launch the scan
         - Write it in DB after scan.
         - Print progress bar
 
-        This function will be called asynchronously.
-        Return one report
+        This function can be called asynchronously.
         """
 
         last_wp_report = self.wp_reports.find(ScanReport(site=wp_site["url"]))
@@ -187,14 +197,17 @@ def _run_scans(self, wp_sites: List[Site]) -> ReportCollection:
         """
 
         log.info(f"Starting scans on {len(wp_sites)} configured sites")
+
+        # reset new reports and scanned sites list. 
+        self._futures.clear()
+        self.new_reports.clear()
+        self.scanner.scanned_sites.clear()
+
         for wp_site in wp_sites:
-            self.futures.append(self.executor.submit(self.scan_site, wp_site))
-        for f in self.futures:
+            self._futures.append(self._executor.submit(self._scan_site, wp_site))
+        for f in self._futures:
             try:
                 self.new_reports.append(f.result())
-            # Handle interruption from inside threads when using --ff
-            except InterruptedError:
-                self.interrupt()
             except concurrent.futures.CancelledError:
                 pass
         # Ensure everything is down
@@ -208,9 +221,6 @@ def run_scans(self) -> Tuple[int, ReportCollection]:
         :Returns: `tuple (exit code, reports)`
         """
 
-        # reset new reports
-        self.new_reports = ReportCollection()
-
         # Check sites are in the config
         if len(self.wp_sites) == 0:
             log.error(
@@ -221,11 +231,14 @@ def run_scans(self) -> Tuple[int, ReportCollection]:
         self.wp_reports.open()
         try:
             self._run_scans(self.wp_sites)
+        # Handle interruption from inside threads when using --ff
+        except InterruptedError:
+            self.interrupt()
         finally:
             self.wp_reports.close()
 
         # Print results and finish
-        self._print_new_reports_results()
+        log.info(repr(self.new_reports))
 
         if not any([r["status"] == "ERROR" for r in self.new_reports if r]):
             log.info("Scans finished successfully.")