admin/Pennyw0rth-NetExec
1
0
Fork 0
mirror of https://github.com/Pennyw0rth/NetExec synced 2026-06-06 16:34:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(database): working on making database operations more async, reducing write concurrency issues, and fixing sqlalchemy query formats

Browse Source
This commit is contained in:
Marshall Hallenbeck
2023-03-05 21:12:13 -05:00
parent f67ebe2154
commit f90871f025
6 changed files with 124 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cme/crackmapexec.py
+24 -13
View File
@@ -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__':
cme/protocols/mssql/database.py
+18 -6
View File
@@ -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
cme/protocols/smb.py
+1 -1
View File
@@ -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:
cme/protocols/smb/database.py
+64 -52
View File
@@ -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 <SessionTransactionState.CLOSED: 5>
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:
poetry.lock
Generated
+16 -1
View File
@@ -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"},
]
pyproject.toml
+1
View File
@@ -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 = "*"