Issue Nekmo#62: Report to file (store lines in UrlsInfo).

michael101096 · Oct 9, 2019 · a31b487 · a31b487
1 parent f22593c
commit a31b487
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/dirhunt/crawler.py b/dirhunt/crawler.py
@@ -23,6 +23,8 @@
 
 
 class Crawler(ThreadPoolExecutor):
+    urls_info = None
+
     def __init__(self, max_workers=None, interesting_extensions=None, interesting_files=None, std=None,
                  progress_enabled=True, timeout=10, depth=3, not_follow_subdomains=False, exclude_sources=(),
                  not_allow_redirects=False, proxies=None, delay=0, limit=1000, to_file=None):
@@ -169,8 +171,8 @@ def print_urls_info(self):
             self.echo(r'No interesting files detected ¯\_(ツ)_/¯')
             return
         self.echo('━' * get_terminal_size()[0])
-        UrlsInfo(self.index_of_processors, self.sessions, self.std, self._max_workers, self.progress_enabled,
-                 self.timeout).start()
+        self.urls_info = UrlsInfo(self.index_of_processors, self.sessions, self.std, self._max_workers,
+                                  self.progress_enabled, self.timeout, bool(self.to_file)).start()
 
     def restart(self):
         try:
@@ -195,5 +197,5 @@ def json(self):
             'index_of_processors': self.index_of_processors,
             'processing': self.processing,
             'processed': self.processed,
-            # TODO: self.results ?
+            # TODO: self.urls_info.lines
         }
diff --git a/dirhunt/url_info.py b/dirhunt/url_info.py
@@ -148,7 +148,8 @@ class UrlsInfo(Pool):
     count = 0
     current = 0
 
-    def __init__(self, processors, sessions, std=None, max_workers=None, progress_enabled=True, timeout=10):
+    def __init__(self, processors, sessions, std=None, max_workers=None, progress_enabled=True, timeout=10,
+                 save_info=False):
         super(UrlsInfo, self).__init__(max_workers)
         self.lock = Lock()
         self.processors = processors
@@ -157,6 +158,8 @@ def __init__(self, processors, sessions, std=None, max_workers=None, progress_en
         self.spinner = random_spinner()
         self.progress_enabled = progress_enabled
         self.timeout = timeout
+        self.lines = []
+        self.save_info = save_info
 
     def callback(self, url_len, extra_len, file):
         line = None
@@ -168,6 +171,8 @@ def callback(self, url_len, extra_len, file):
             self.error_files += 1
         self.lock.acquire()
         self.erase()
+        if self.save_info:
+            self.lines.append(line)
         if line:
             self.echo(line)
         self.print_progress()