if script fails, report to telegram too

1 files changed, 52 insertions, 41 deletions

diff --git a/suddenly-opened-ports-checker b/suddenly-opened-ports-checker
index b4c1972..8e14845 100755
--- a/suddenly-opened-ports-checker
+++ b/suddenly-opened-ports-checker
@@ -2,10 +2,12 @@
 import logging
 import yaml
 import math
+import html
 
 from argparse import ArgumentParser
 from lib.worker import Worker
 from lib.results import Results
+from ch1p import telegram_notify
 
 logger = logging.getLogger(__name__)
 
@@ -40,47 +42,56 @@ def main():
     if not args.no_telegram:
         assert 'telegram' in config
 
-    # let's go
-    results = Results()
-    max_threads = math.inf if args.threads_limit == 0 else args.threads_limit
-    active_threads = 1
-
-    def get_active_threads():
-        n = active_threads
-        if workers:
-            n += workers[0].concurrency
-        return n
-
-    workers = []
-    for name, data in config['servers'].items():
-        w = Worker(name, data['host'], data['opened'], data['ignore'],
-                   concurrency=int(data['concurrency']) if 'concurrency' in data else args.concurrency,
-                   timeout=int(data['timeout']) if 'timeout' in data else args.timeout)
-        workers.append(w)
-
-    current_workers = []
-    while workers:
-        w = workers.pop(0)
-        active_threads += w.concurrency+1
-
-        current_workers.append(w)
-        w.start()
-
-        while current_workers and get_active_threads() >= max_threads:
-            for cw in current_workers:
-                cw.join(timeout=0.1)
-                if not cw.is_alive():
-                    results.add(cw)
-                    current_workers.remove(cw)
-                    active_threads -= cw.concurrency+1
-
-    for cw in current_workers:
-        cw.join()
-        results.add(cw)
-
-    if results.has_warnings() and not args.no_telegram:
-        results.notify(chat_id=config['telegram']['chat-id'],
-                       token=config['telegram']['token'])
+    try:
+        # let's go
+        results = Results()
+        max_threads = math.inf if args.threads_limit == 0 else args.threads_limit
+        active_threads = 1
+
+        def get_active_threads():
+            n = active_threads
+            if workers:
+                n += workers[0].concurrency
+            return n
+
+        workers = []
+        for name, data in config['servers'].items():
+            w = Worker(name, data['host'], data['opened'], data['ignore'],
+                       concurrency=int(data['concurrency']) if 'concurrency' in data else args.concurrency,
+                       timeout=int(data['timeout']) if 'timeout' in data else args.timeout)
+            workers.append(w)
+
+        current_workers = []
+        while workers:
+            w = workers.pop(0)
+            active_threads += w.concurrency+1
+
+            current_workers.append(w)
+            w.start()
+
+            while current_workers and get_active_threads() >= max_threads:
+                for cw in current_workers:
+                    cw.join(timeout=0.1)
+                    if not cw.is_alive():
+                        results.add(cw)
+                        current_workers.remove(cw)
+                        active_threads -= cw.concurrency+1
+
+        for cw in current_workers:
+            cw.join()
+            results.add(cw)
+
+        if results.has_warnings() and not args.no_telegram:
+            results.notify(chat_id=config['telegram']['chat-id'],
+                           token=config['telegram']['token'])
+    except KeyboardInterrupt:
+        pass
+    except Exception as e:
+        if not args.no_telegram:
+            telegram_notify(html.escape(str(e)),
+                            parse_mode='html',
+                            chat_id=config['telegram']['chat-id'],
+                            token=config['telegram']['token'])
 
 
 if __name__ == '__main__':