Misc bugfixes, asyncio speed optimizations

bbot/cli.py

@@ -6,9 +6,9 @@
 import asyncio
 import logging
 import traceback
-from aioconsole import ainput
 from omegaconf import OmegaConf
 from contextlib import suppress
+from aioconsole import stream
 
 # fix tee buffering
 sys.stdout.reconfigure(line_buffering=True)
@@ -20,6 +20,7 @@
 from bbot import __version__
 from bbot.modules import module_loader
 from bbot.core.configurator.args import parser
+from bbot.core.helpers.misc import smart_decode
 from bbot.core.helpers.logger import log_to_stderr
 from bbot.core.configurator import ensure_config_files, check_cli_args, environ
 
@@ -301,46 +302,56 @@ async def _main():
 
                 if not options.dry_run:
                     log.trace(f"Command: {' '.join(sys.argv)}")
-                    if not options.agent_mode and not options.yes and sys.stdin.isatty():
-                        log.hugesuccess(f"Scan ready. Press enter to execute {scanner.name}")
-                        input()
-
-                    def handle_keyboard_input(keyboard_input):
-                        kill_regex = re.compile(r"kill (?P<module>[a-z0-9_]+)")
-                        if keyboard_input:
-                            log.verbose(f'Got keyboard input: "{keyboard_input}"')
-                            kill_match = kill_regex.match(keyboard_input)
-                            if kill_match:
-                                module = kill_match.group("module")
-                                if module in scanner.modules:
-                                    log.hugewarning(f'Killing module: "{module}"')
-                                    scanner.manager.kill_module(module, message="killed by user")
-                                else:
-                                    log.warning(f'Invalid module: "{module}"')
-                        else:
-                            toggle_log_level(logger=log)
-                            scanner.manager.modules_status(_log=True)
 
-                    async def akeyboard_listen():
-                        allowed_errors = 10
-                        while 1:
-                            keyboard_input = "a"
+                    if sys.stdin.isatty():
+                        if not options.agent_mode and not options.yes:
+                            log.hugesuccess(f"Scan ready. Press enter to execute {scanner.name}")
+                            input()
+
+                        def handle_keyboard_input(keyboard_input):
+                            kill_regex = re.compile(r"kill (?P<module>[a-z0-9_]+)")
+                            if keyboard_input:
+                                log.verbose(f'Got keyboard input: "{keyboard_input}"')
+                                kill_match = kill_regex.match(keyboard_input)
+                                if kill_match:
+                                    module = kill_match.group("module")
+                                    if module in scanner.modules:
+                                        log.hugewarning(f'Killing module: "{module}"')
+                                        scanner.manager.kill_module(module, message="killed by user")
+                                    else:
+                                        log.warning(f'Invalid module: "{module}"')
+                            else:
+                                toggle_log_level(logger=log)
+                                scanner.manager.modules_status(_log=True)
+
+                        # Reader
+                        reader = stream.StandardStreamReader()
+                        protocol = stream.StandardStreamReaderProtocol(reader)
+                        await asyncio.get_event_loop().connect_read_pipe(lambda: protocol, sys.stdin)
+
+                        async def akeyboard_listen():
                             try:
-                                keyboard_input = await ainput()
-                            except Exception:
-                                allowed_errors -= 1
-                            handle_keyboard_input(keyboard_input)
-                            if allowed_errors <= 0:
-                                break
-
-                    try:
-                        keyboard_listen_task = asyncio.create_task(akeyboard_listen())
-
-                        await scanner.async_start_without_generator()
-                    finally:
-                        keyboard_listen_task.cancel()
-                        with suppress(asyncio.CancelledError):
-                            await keyboard_listen_task
+                                allowed_errors = 10
+                                while 1:
+                                    keyboard_input = None
+                                    try:
+                                        keyboard_input = smart_decode((await reader.readline()).strip())
+                                        allowed_errors = 10
+                                    except Exception as e:
+                                        log_to_stderr(f"Error in keyboard listen loop: {e}", level="TRACE")
+                                        log_to_stderr(traceback.format_exc(), level="TRACE")
+                                        allowed_errors -= 1
+                                    if keyboard_input is not None:
+                                        handle_keyboard_input(keyboard_input)
+                                    if allowed_errors <= 0:
+                                        break
+                            except Exception as e:
+                                log_to_stderr(f"Error in keyboard listen task: {e}", level="ERROR")
+                                log_to_stderr(traceback.format_exc(), level="TRACE")
+
+                        asyncio.create_task(akeyboard_listen())
+
+                    await scanner.async_start_without_generator()
 
             except bbot.core.errors.ScanError as e:
                 log_to_stderr(str(e), level="ERROR")

bbot/core/helpers/async_helpers.py

@@ -5,13 +5,12 @@
 import threading
 from datetime import datetime
 from queue import Queue, Empty
+from cachetools import LRUCache
 from .misc import human_timedelta
 from contextlib import asynccontextmanager
 
 log = logging.getLogger("bbot.core.helpers.async_helpers")
 
-from .cache import CacheDict
-
 
 class ShuffleQueue(asyncio.Queue):
     def _put(self, item):
@@ -32,20 +31,20 @@ class NamedLock:
     """
     Returns a unique asyncio.Lock() based on a provided string
 
-    Useful for preventing multiple operations from occuring on the same data in parallel
+    Useful for preventing multiple operations from occurring on the same data in parallel
     E.g. simultaneous DNS lookups on the same hostname
     """
 
     def __init__(self, max_size=1000):
-        self._cache = CacheDict(max_size=max_size)
+        self._cache = LRUCache(maxsize=max_size)
 
     @asynccontextmanager
     async def lock(self, name):
         try:
-            lock = self._cache.get(name)
+            lock = self._cache[name]
         except KeyError:
             lock = _Lock(name)
-            self._cache.put(name, lock)
+            self._cache[name] = lock
         async with lock:
             yield
 

bbot/core/helpers/cache.py

@@ -1,8 +1,6 @@
 import os
 import time
 import logging
-from contextlib import suppress
-from collections import OrderedDict
 
 from .misc import sha1
 
@@ -53,84 +51,3 @@ def is_cached(self, key, cache_hrs=24 * 7):
 
 def cache_filename(self, key):
     return self.cache_dir / sha1(key).hexdigest()
-
-
-_sentinel = object()
-
-
-class CacheDict:
-    """
-    Dictionary to store cached values, with a maximum size limit
-    """
-
-    def __init__(self, max_size=1000):
-        self._cache = OrderedDict()
-        self._max_size = int(max_size)
-
-    def get(self, name, fallback=_sentinel):
-        name_hash = self._hash(name)
-        try:
-            return self._cache[name_hash]
-        except KeyError:
-            if fallback is not _sentinel:
-                return fallback
-            raise
-        finally:
-            with suppress(KeyError):
-                self._cache.move_to_end(name_hash)
-            self._truncate()
-
-    def put(self, name, value):
-        name_hash = self._hash(name)
-        try:
-            self._cache[name_hash] = value
-        finally:
-            with suppress(KeyError):
-                self._cache.move_to_end(name_hash)
-            self._truncate()
-
-    def _truncate(self):
-        if not self or len(self) <= self._max_size:
-            return
-        for nh in list(self._cache.keys()):
-            try:
-                del self._cache[nh]
-            except KeyError:
-                pass
-            if not self or len(self) <= self._max_size:
-                break
-
-    def keys(self):
-        return self._cache.keys()
-
-    def values(self):
-        return self._cache.values()
-
-    def items(self):
-        return self._cache.items()
-
-    def clear(self):
-        return self._cache.clear()
-
-    def _hash(self, v):
-        if type(v) == int:
-            return v
-        return hash(str(v))
-
-    def __contains__(self, item):
-        return self._hash(item) in self._cache
-
-    def __iter__(self):
-        return iter(self._cache)
-
-    def __getitem__(self, item):
-        return self.get(item)
-
-    def __setitem__(self, item, value):
-        self.put(item, value)
-
-    def __bool__(self):
-        return bool(self._cache)
-
-    def __len__(self):
-        return len(self._cache)

bbot/core/helpers/depsinstaller/installer.py

@@ -157,9 +157,9 @@ async def pip_install(self, packages, constraints=None):
         command = [sys.executable, "-m", "pip", "install", "--upgrade"] + packages
 
         if constraints:
-            contraints_tempfile = self.parent_helper.tempfile(constraints, pipe=False)
+            constraints_tempfile = self.parent_helper.tempfile(constraints, pipe=False)
             command.append("--constraint")
-            command.append(contraints_tempfile)
+            command.append(constraints_tempfile)
 
         process = None
         try: