Skip to content

Commit

Permalink
Merge pull request #1051 from blacklanternsecurity/speed-optimizations
Browse files Browse the repository at this point in the history 
Misc bugfixes, asyncio speed optimizations
  • Loading branch information
@TheTechromancer
TheTechromancer authored Feb 14, 2024
2 parents e952e7f + 28f24ac commit d506409
Show file tree
Hide file tree
Showing 32 changed files with 485 additions and 400 deletions.
89 changes: 50 additions & 39 deletions bbot/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

Expand Down Expand Up @@ -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")
Expand Down
11 changes: 5 additions & 6 deletions bbot/core/helpers/async_helpers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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

Expand Down
83 changes: 0 additions & 83 deletions bbot/core/helpers/cache.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
import os
import time
import logging
from contextlib import suppress
from collections import OrderedDict

from .misc import sha1

Expand Down Expand Up @@ -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)
4 changes: 2 additions & 2 deletions bbot/core/helpers/depsinstaller/installer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading

0 comments on commit d506409

Please sign in to comment.