diff --git a/cme/crackmapexec.py b/cme/crackmapexec.py index adc34d95..f7cf182c 100755 --- a/cme/crackmapexec.py +++ b/cme/crackmapexec.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- +from sqlalchemy.exc import SAWarning from cme.logger import setup_logger, setup_debug_logger, CMEAdapter from cme.helpers.logger import highlight @@ -30,19 +31,28 @@ import random import os import sys import logging -from sqlalchemy import create_engine, MetaData - -from sqlalchemy.ext.declarative import DeferredReflection from sqlalchemy.orm import declarative_base +from sqlalchemy.ext.asyncio import create_async_engine +import warnings Base = declarative_base() setup_logger() logger = CMEAdapter() +# if there is an issue with SQLAlchemy and a connection cannot be cleaned up properly it spews out annoying warnings +warnings.filterwarnings("ignore", category=SAWarning) -class Computers(DeferredReflection, Base): - __tablename__ = "computers" + +def create_db_engine(db_path): + db_engine = create_async_engine( + f"sqlite+aiosqlite:///{db_path}", + isolation_level="AUTOCOMMIT", + future=True + ) # can add echo=True + # db_engine.execution_options(isolation_level="AUTOCOMMIT") + # db_engine.connect().connection.text_factory = str + return db_engine async def monitor_threadpool(pool, targets): @@ -86,7 +96,7 @@ async def run_protocol(loop, protocol_obj, args, db, target, jitter): except asyncio.TimeoutError: logging.debug("Thread exceeded timeout") except asyncio.CancelledError: - logging.debug("Stopping thread") + logging.debug("Shutting down DB") thread.cancel() except sqlite3.OperationalError as e: logging.debug("Sqlite error - sqlite3.operationalError - {}".format(str(e))) @@ -121,11 +131,14 @@ async def start_threadpool(protocol_obj, args, db, targets, jitter): logger.info("Shutting down, please wait...") logging.debug("Cancelling scan") finally: + await asyncio.shield(db.shutdown_db()) monitor_task.cancel() pool.shutdown(wait=True) def main(): + logging.getLogger('aiosqlite').setLevel(logging.CRITICAL) + logging.getLogger('sqlalchemy.pool.impl.NullPool').setLevel(logging.CRITICAL) first_run_setup(logger) args = gen_cli_args() @@ -203,19 +216,16 @@ def main(): logging.debug(f"Protocol DB Path: {protocol_db_path}") protocol_object = getattr(p_loader.load_protocol(protocol_path), args.protocol) + logging.debug(f"Protocol Object: {protocol_object}") protocol_db_object = getattr(p_loader.load_protocol(protocol_db_path), 'database') + logging.debug(f"Protocol DB Object: {protocol_object}") db_path = os.path.join(CME_PATH, 'workspaces', current_workspace, args.protocol + '.db') logging.debug(f"DB Path: {db_path}") - db_engine = create_engine(f"sqlite:///{db_path}") - db_engine.execution_options(isolation_level="AUTOCOMMIT") - db_engine.connect().connection.text_factory = str + db_engine = create_db_engine(db_path) - metadata = MetaData() - metadata.reflect(bind=db_engine) - - db = protocol_db_object(db_engine, metadata=metadata) + db = protocol_db_object(db_engine) setattr(protocol_object, 'config', config) @@ -280,6 +290,7 @@ def main(): finally: if module_server: module_server.shutdown() + asyncio.run(db_engine.dispose()) if __name__ == '__main__': diff --git a/cme/protocols/mssql/database.py b/cme/protocols/mssql/database.py index 7d56adc0..93525bf9 100755 --- a/cme/protocols/mssql/database.py +++ b/cme/protocols/mssql/database.py @@ -1,15 +1,20 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -from sqlalchemy.orm import sessionmaker -from sqlalchemy import func import logging +from sqlalchemy import MetaData, func +from sqlalchemy.orm import sessionmaker, scoped_session +from sqlalchemy.ext.asyncio import AsyncSession class database: def __init__(self, db_engine, metadata=None): - session = sessionmaker(bind=db_engine) + metadata = MetaData() + metadata.reflect(bind=db_engine) + + session_factory = sessionmaker(bind=db_engine, expire_on_commit=False, class_=AsyncSession) + Session = scoped_session(session_factory) # this is still named "conn" when it is the session object; TODO: rename - self.conn = session() + self.conn = Session() self.metadata = metadata self.computers_table = metadata.tables["computers"] self.admin_relations_table = metadata.tables["admin_relations"] @@ -86,7 +91,8 @@ class database: [new_host] ) except Exception as e: - logging.error(f"Exception: {e}") + pass + # logging.error(f"Exception: {e}") else: for host in results: try: @@ -98,7 +104,13 @@ class database: ) ) except Exception as e: - logging.error(f"Exception: {e}") + pass + # logging.error(f"Exception: {e}") + try: + self.conn.commit() + except Exception as e: + logging.error(f"Exception while committing to database: {e}") + self.conn.close() return cid diff --git a/cme/protocols/smb.py b/cme/protocols/smb.py index 805f1d6b..3f72fc76 100755 --- a/cme/protocols/smb.py +++ b/cme/protocols/smb.py @@ -279,7 +279,7 @@ class smb(connection): ''' self.conn.logoff() except Exception as e: - logging.debug(e) + logging.debug(f"Error logging off system: {e}") pass if self.args.domain: diff --git a/cme/protocols/smb/database.py b/cme/protocols/smb/database.py index c5e7b3d4..caa20aed 100755 --- a/cme/protocols/smb/database.py +++ b/cme/protocols/smb/database.py @@ -3,11 +3,11 @@ import logging from sqlalchemy import MetaData, func, inspect, Table, select, insert, update, delete from sqlalchemy.dialects.sqlite import insert as sqlite_upsert +from sqlalchemy.exc import IllegalStateChangeError from sqlalchemy.orm import sessionmaker, scoped_session from sqlalchemy.ext.asyncio import AsyncSession from datetime import datetime import asyncio -import copy def get_table_names(conn): @@ -28,12 +28,21 @@ class database: self.db_engine = db_engine self.metadata = MetaData() asyncio.run(self.reflect_tables()) - session_factory = sessionmaker(bind=self.db_engine, expire_on_commit=False, class_=AsyncSession) + session_factory = sessionmaker(bind=self.db_engine, expire_on_commit=True, class_=AsyncSession) # session_factory = sessionmaker(bind=self.db_engine, expire_on_commit=False) Session = scoped_session(session_factory) # this is still named "conn" when it is the session object; TODO: rename self.conn = Session() + async def shutdown_db(self): + try: + await asyncio.shield(self.conn.close()) + # due to the async nature of CME, sometimes session state is a bit messy and this will throw: + # Method 'close()' can't be called here; method '_connection_for_bind()' is already in progress and + # this would cause an unexpected state change to + except IllegalStateChangeError as e: + logging.debug(f"Error while closing session db object: {e}") + async def reflect_tables(self): async with self.db_engine.connect() as conn: await conn.run_sync(self.metadata.reflect) @@ -233,7 +242,6 @@ class database: """ Check if this host has already been added to the database, if not add it in. """ - logging.debug(f"Inside add_computer") domain = domain.split('.')[0].upper() update_hosts = [] @@ -244,37 +252,29 @@ class database: results = res.all() logging.debug(f"Results in add_computer: {results}") + host = { + "ip": ip, + "hostname": hostname, + "domain": domain, + "os": os, + "dc": dc, + "smbv1": smbv1, + "signing": signing, + "spooler": spooler, + "zerologon": zerologon, + "petitpotam": petitpotam + } + + # create new computer if not results: - logging.debug(f"Results apparently empty") - new_host = { - "ip": ip, - "hostname": hostname, - "domain": domain, - "os": os, - "dc": dc, - "smbv1": smbv1, - "signing": signing, - "spooler": spooler, - "zerologon": zerologon, - "petitpotam": petitpotam - } - try: - cid = asyncio.run(self.conn.execute( - insert(self.ComputersTable).values(new_host).returning(self.ComputersTable.c.id) - )).scalar_one() - except Exception as e: - logging.error(f"Exception: {e}") + update_hosts = [host] + # update existing hosts data else: for host in results: - print(type(host)) - print(host) - print(dir(host)) - #row_as_dict = {column: str(getattr(host, column)) for column in host.__table__.c.keys()} - #print(row_as_dict) - computer_data = {"id": host.id} + computer_data = host._asdict() # only update column if it is being passed in if ip is not None: - computer_data["ip"] = ip + computer_data["ip"] = ip if hostname is not None: computer_data["hostname"] = hostname if domain is not None: @@ -294,16 +294,15 @@ class database: if dc is not None: computer_data["dc"] = dc update_hosts.append(computer_data) - print(f"Update Hosts: {update_hosts}") - cid = asyncio.run( - self.conn.execute( - sqlite_upsert(self.ComputersTable), - update_hosts - ) + + logging.debug(f"Update Hosts: {update_hosts}") + asyncio.run( + self.conn.execute( + sqlite_upsert(self.ComputersTable), + update_hosts ) - logging.debug(f"CID: {cid}") - logging.debug(f"CID Type: {type(cid)}") - return cid + ) + # asyncio.run(self.conn.close()) def add_credential(self, credtype, domain, username, password, group_id=None, pillaged_from=None): """ @@ -533,21 +532,27 @@ class database: domain = domain.split('.')[0].upper() if user_id: - users = self.conn.query(self.UsersTable).filter( + q = select(self.UsersTable).filter( self.UsersTable.c.id == user_id - ).all() + ) + res = asyncio.run(self.conn.execute(q)) + users = res.all() else: - users = self.conn.query(self.UsersTable).filter( + q = select(self.UsersTable).filter( self.UsersTable.c.credtype == credtype, func.lower(self.UsersTable.c.domain) == func.lower(domain), func.lower(self.UsersTable.c.username) == func.lower(username), self.UsersTable.c.password == password - ).all() + ) + res = asyncio.run(self.conn.execute(q)) + users = res.all() logging.debug(f"Users: {users}") - hosts = self.conn.query(self.ComputersTable).filter( + q = select(self.ComputersTable).filter( self.ComputersTable.c.ip.like(func.lower(f"%{host}%")) ) + res = asyncio.run(self.conn.execute(q)) + hosts = res.all() logging.debug(f"Hosts: {hosts}") if users is not None and hosts is not None: @@ -556,19 +561,26 @@ class database: host_id = host[0] # Check to see if we already added this link - links = self.conn.query(self.AdminRelationsTable).filter( + # links = self.conn.query(self.AdminRelationsTable).filter( + # self.AdminRelationsTable.c.userid == user_id, + # self.AdminRelationsTable.c.computerid == host_id + # ).all() + q = select(self.AdminRelationsTable).filter( self.AdminRelationsTable.c.userid == user_id, self.AdminRelationsTable.c.computerid == host_id - ).all() + ) + res = asyncio.run(self.conn.execute(q)) + links = res.all() if not links: - self.conn.execute( - self.AdminRelationsTable.insert(), - [{"userid": user_id, "computerid": host_id}] - ) - - self.conn.commit() - self.conn.close() + link = {"userid": user_id, "computerid": host_id} + # self.conn.execute( + # self.AdminRelationsTable.insert(), + # [{"userid": user_id, "computerid": host_id}] + # ) + asyncio.run(self.conn.execute( + insert(self.AdminRelationsTable).values(link) + )) def get_admin_relations(self, user_id=None, host_id=None): if user_id: diff --git a/poetry.lock b/poetry.lock index 4beff607..84caf890 100644 --- a/poetry.lock +++ b/poetry.lock @@ -62,6 +62,17 @@ unicrypto = ">=0.0.9" wcwidth = "*" winacl = ">=0.1.5" +[[package]] +name = "aiosqlite" +version = "0.18.0" +description = "asyncio bridge to the standard sqlite3 module" +category = "main" +optional = false +python-versions = ">=3.7" + +[package.dependencies] +typing_extensions = {version = ">=4.0", markers = "python_version < \"3.8\""} + [[package]] name = "aiowinreg" version = "0.0.9" @@ -1223,7 +1234,7 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more [metadata] lock-version = "1.1" python-versions = "^3.7.0" -content-hash = "891ef2102b8ef5e04fb377170dab86aca822ab420f1cb1ac8ba552dc471535c9" +content-hash = "60f4adf6040782e03179b633b36096915292191bf761639a1834c2903ae4012f" [metadata.files] aardwolf = [ @@ -1247,6 +1258,10 @@ aiosmb = [ {file = "aiosmb-0.4.4-py3-none-any.whl", hash = "sha256:ed70967243ec6634ac443de4e048c666f372ec863e49e8a2c1e01d640989f95e"}, {file = "aiosmb-0.4.4.tar.gz", hash = "sha256:20620498cf5e6794fea29ddcb46afbda8cbf714e3cf8e1da149699f03453637f"}, ] +aiosqlite = [ + {file = "aiosqlite-0.18.0-py3-none-any.whl", hash = "sha256:c3511b841e3a2c5614900ba1d179f366826857586f78abd75e7cbeb88e75a557"}, + {file = "aiosqlite-0.18.0.tar.gz", hash = "sha256:faa843ef5fb08bafe9a9b3859012d3d9d6f77ce3637899de20606b7fc39aa213"}, +] aiowinreg = [ {file = "aiowinreg-0.0.9-py3-none-any.whl", hash = "sha256:8fd39c039021296d47c023f4db863bf6016882c87a50de9870c3471b00ddc148"}, ] diff --git a/pyproject.toml b/pyproject.toml index 91b72843..4ae8de1f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ minikerberos = "0.3.5" aardwolf = "0.2.5" masky = "^0.2.0" sqlalchemy = "^2.0.4" +aiosqlite = "^0.18.0" [tool.poetry.dev-dependencies] flake8 = "*"