From cf375e47bb809043233e486910711215e7da64e1 Mon Sep 17 00:00:00 2001 From: Marshall Hallenbeck Date: Mon, 20 Mar 2023 21:06:21 -0400 Subject: [PATCH 1/6] tests(smb): add setup, teardown, initial add_host test, and definitions for all smb function tests --- poetry.lock | 0 tests/test_smb_database.py | 152 +++++++++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) delete mode 100644 poetry.lock create mode 100644 tests/test_smb_database.py diff --git a/poetry.lock b/poetry.lock deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/test_smb_database.py b/tests/test_smb_database.py new file mode 100644 index 00000000..9372e85d --- /dev/null +++ b/tests/test_smb_database.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +import asyncio +import os +from time import sleep + +from cme.cmedb import CMEDBMenu, create_workspace, create_db_engine, delete_workspace +from cme.first_run import first_run_setup +from cme.loaders.protocol_loader import protocol_loader +from cme.logger import setup_logger, CMEAdapter +from cme.paths import CONFIG_PATH, WS_PATH + + +class TestSmbDatabase: + def setup_class(self): + proto = "smb" + setup_logger() + logger = CMEAdapter() + first_run_setup(logger) + p_loader = protocol_loader() + protocols = p_loader.get_protocols() + create_workspace("test", p_loader, protocols) + + protocol_db_path = p_loader.get_protocols()[proto]["dbpath"] + protocol_db_object = getattr(p_loader.load_protocol(protocol_db_path), "database") + db_path = os.path.join(WS_PATH, "test/smb.db") + db_engine = create_db_engine(db_path) + self.db = protocol_db_object(db_engine) + + def teardown_class(self): + asyncio.run(self.db.shutdown_db()) + delete_workspace("test") + + def test_add_host(self): + print(self.db) + self.db.add_host( + "127.0.0.1", + "localhost", + "TEST.DEV" + "Windows Testing 2023", + True, + False, + True, + False, + True + ) + + def test_update_host(self): + pass + + def test_add_credential(self): + pass + + def test_update_credential(self): + pass + + def test_remove_credential(self): + pass + + def test_add_admin_user(self): + pass + + def test_get_admin_relations(self): + pass + + def test_remove_admin_relation(self): + pass + + def test_is_credential_valid(self): + pass + + def test_get_credentials(self): + pass + + def test_get_credential(self): + pass + + def test_is_credential_local(self): + pass + + def test_is_host_valid(self): + pass + + def test_get_hosts(self): + pass + + def test_is_group_valid(self): + pass + + def test_add_group(self): + pass + + def test_get_groups(self): + pass + + def test_get_group_relations(self): + pass + + def test_remove_group_relations(self): + pass + + def test_is_user_valid(self): + pass + + def test_get_users(self): + pass + + def test_get_user(self): + pass + + def test_get_domain_controllers(self): + pass + + def test_is_share_valid(self): + pass + + def test_add_share(self): + pass + + def test_get_shares(self): + pass + + def test_get_shares_by_access(self): + pass + + def test_get_users_with_share_access(self): + pass + + def test_add_domain_backupkey(self): + pass + + def test_get_domain_backupkey(self): + pass + + def test_is_dpapi_secret_valid(self): + pass + + def test_add_dpapi_secrets(self): + pass + + def test_get_dpapi_secrets(self): + pass + + def test_add_loggedin_relation(self): + pass + + def test_get_loggedin_relations(self): + pass + + def test_remove_loggedin_relations(self): + pass + From 470e9865602fff6f8d9456594891914686e07bb4 Mon Sep 17 00:00:00 2001 From: Marshall Hallenbeck Date: Mon, 20 Mar 2023 21:13:15 -0400 Subject: [PATCH 2/6] tests(deps): add pytest dependency --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index ca54abb9..c6d0ebe8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ asyauth = "^0.0.12" masky = "^0.2.0" sqlalchemy = "^2.0.4" aiosqlite = "^0.18.0" +pytest = "^7.2.2" [tool.poetry.dev-dependencies] flake8 = "*" From 22a2e281e6de4761c1a8e93b733deaaf16ae931c Mon Sep 17 00:00:00 2001 From: Marshall Hallenbeck Date: Mon, 20 Mar 2023 21:14:07 -0400 Subject: [PATCH 3/6] update how workspaces are created so tests can utilize functionality --- cme/cmedb.py | 65 ++++++++++++++++++++++++++------------------- cme/crackmapexec.py | 2 +- cme/paths.py | 3 ++- 3 files changed, 40 insertions(+), 30 deletions(-) diff --git a/cme/cmedb.py b/cme/cmedb.py index cf6402ed..77cb54e6 100644 --- a/cme/cmedb.py +++ b/cme/cmedb.py @@ -3,6 +3,7 @@ import cmd import logging +import shutil import sqlite3 import sys import os @@ -10,7 +11,7 @@ import requests from terminaltables import AsciiTable import configparser from cme.loaders.protocol_loader import protocol_loader -from cme.paths import CONFIG_PATH, WS_PATH +from cme.paths import CONFIG_PATH, WS_PATH, WORKSPACE_DIR from requests import ConnectionError from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.exc import SAWarning @@ -275,7 +276,7 @@ class CMEDBMenu(cmd.Cmd): print("[-] Error reading cme.conf: {}".format(e)) sys.exit(1) - self.workspace_dir = os.path.expanduser('~/.cme/workspaces') + self.conn = None self.p_loader = protocol_loader() self.protocols = self.p_loader.get_protocols() @@ -295,7 +296,7 @@ class CMEDBMenu(cmd.Cmd): if not proto: return - proto_db_path = os.path.join(self.workspace_dir, self.workspace, proto + '.db') + proto_db_path = os.path.join(WORKSPACE_DIR, self.workspace, proto + '.db') if os.path.exists(proto_db_path): self.conn = create_db_engine(proto_db_path) db_nav_object = self.p_loader.load_protocol(self.protocols[proto]['nvpath']) @@ -326,38 +327,16 @@ class CMEDBMenu(cmd.Cmd): if subcommand == 'create': new_workspace = line.split()[1].strip() print("[*] Creating workspace '{}'".format(new_workspace)) - os.mkdir(os.path.join(self.workspace_dir, new_workspace)) - - for protocol in self.protocols.keys(): - try: - protocol_object = self.p_loader.load_protocol(self.protocols[protocol]['dbpath']) - except KeyError: - continue - proto_db_path = os.path.join(self.workspace_dir, new_workspace, protocol + '.db') - - if not os.path.exists(proto_db_path): - print('[*] Initializing {} protocol database'.format(protocol.upper())) - conn = sqlite3.connect(proto_db_path) - c = conn.cursor() - - # try to prevent some weird sqlite I/O errors - c.execute('PRAGMA journal_mode = OFF') - c.execute('PRAGMA foreign_keys = 1') - - getattr(protocol_object, 'database').db_schema(c) - - # commit the changes and close everything off - conn.commit() - conn.close() + self.create_workspace(new_workspace, self.p_loader, self.protocols) self.do_workspace(new_workspace) elif subcommand == 'list': print("[*] Enumerating Workspaces") - for workspace in os.listdir(os.path.join(self.workspace_dir)): + for workspace in os.listdir(os.path.join(WORKSPACE_DIR)): if workspace == self.workspace: print("==> "+workspace) else: print(workspace) - elif os.path.exists(os.path.join(self.workspace_dir, line)): + elif os.path.exists(os.path.join(WORKSPACE_DIR, line)): self.config.set('CME', 'workspace', line) self.write_configfile() self.workspace = line @@ -370,6 +349,36 @@ class CMEDBMenu(cmd.Cmd): print_help(help_string) +def create_workspace(workspace_name, p_loader, protocols): + os.mkdir(os.path.join(WORKSPACE_DIR, workspace_name)) + + for protocol in protocols.keys(): + try: + protocol_object = p_loader.load_protocol(protocols[protocol]['dbpath']) + except KeyError: + continue + proto_db_path = os.path.join(WORKSPACE_DIR, workspace_name, protocol + '.db') + + if not os.path.exists(proto_db_path): + print('[*] Initializing {} protocol database'.format(protocol.upper())) + conn = sqlite3.connect(proto_db_path) + c = conn.cursor() + + # try to prevent some weird sqlite I/O errors + c.execute('PRAGMA journal_mode = OFF') + c.execute('PRAGMA foreign_keys = 1') + + getattr(protocol_object, 'database').db_schema(c) + + # commit the changes and close everything off + conn.commit() + conn.close() + + +def delete_workspace(workspace_name): + shutil.rmtree(os.path.join(WORKSPACE_DIR, workspace_name)) + + def initialize_db(logger): if not os.path.exists(os.path.join(WS_PATH, 'default')): logger.info('Creating default workspace') diff --git a/cme/crackmapexec.py b/cme/crackmapexec.py index fd986c1b..10b71e48 100755 --- a/cme/crackmapexec.py +++ b/cme/crackmapexec.py @@ -228,7 +228,7 @@ def main(): 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}") + logging.debug(f"Protocol DB Object: {protocol_db_object}") db_path = os.path.join(CME_PATH, 'workspaces', current_workspace, args.protocol + '.db') logging.debug(f"DB Path: {db_path}") diff --git a/cme/paths.py b/cme/paths.py index 8c4a0eff..795bb088 100644 --- a/cme/paths.py +++ b/cme/paths.py @@ -9,4 +9,5 @@ if hasattr(sys, 'getandroidapilevel'): TMP_PATH = os.path.join('/data','data', 'com.termux', 'files', 'usr', 'tmp', 'cme_hosted') WS_PATH = os.path.join(CME_PATH, 'workspaces') CERT_PATH = os.path.join(CME_PATH, 'cme.pem') -CONFIG_PATH = os.path.join(CME_PATH, 'cme.conf') \ No newline at end of file +CONFIG_PATH = os.path.join(CME_PATH, 'cme.conf') +WORKSPACE_DIR = os.path.join(CME_PATH, 'workspaces') From f063c711cc9d932f32f0daa5299d9405ec0df191 Mon Sep 17 00:00:00 2001 From: Marshall Hallenbeck Date: Mon, 20 Mar 2023 21:06:21 -0400 Subject: [PATCH 4/6] tests(smb): add setup, teardown, initial add_host test, and definitions for all smb function tests --- poetry.lock | 0 tests/test_smb_database.py | 152 +++++++++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) delete mode 100644 poetry.lock create mode 100644 tests/test_smb_database.py diff --git a/poetry.lock b/poetry.lock deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/test_smb_database.py b/tests/test_smb_database.py new file mode 100644 index 00000000..9372e85d --- /dev/null +++ b/tests/test_smb_database.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +import asyncio +import os +from time import sleep + +from cme.cmedb import CMEDBMenu, create_workspace, create_db_engine, delete_workspace +from cme.first_run import first_run_setup +from cme.loaders.protocol_loader import protocol_loader +from cme.logger import setup_logger, CMEAdapter +from cme.paths import CONFIG_PATH, WS_PATH + + +class TestSmbDatabase: + def setup_class(self): + proto = "smb" + setup_logger() + logger = CMEAdapter() + first_run_setup(logger) + p_loader = protocol_loader() + protocols = p_loader.get_protocols() + create_workspace("test", p_loader, protocols) + + protocol_db_path = p_loader.get_protocols()[proto]["dbpath"] + protocol_db_object = getattr(p_loader.load_protocol(protocol_db_path), "database") + db_path = os.path.join(WS_PATH, "test/smb.db") + db_engine = create_db_engine(db_path) + self.db = protocol_db_object(db_engine) + + def teardown_class(self): + asyncio.run(self.db.shutdown_db()) + delete_workspace("test") + + def test_add_host(self): + print(self.db) + self.db.add_host( + "127.0.0.1", + "localhost", + "TEST.DEV" + "Windows Testing 2023", + True, + False, + True, + False, + True + ) + + def test_update_host(self): + pass + + def test_add_credential(self): + pass + + def test_update_credential(self): + pass + + def test_remove_credential(self): + pass + + def test_add_admin_user(self): + pass + + def test_get_admin_relations(self): + pass + + def test_remove_admin_relation(self): + pass + + def test_is_credential_valid(self): + pass + + def test_get_credentials(self): + pass + + def test_get_credential(self): + pass + + def test_is_credential_local(self): + pass + + def test_is_host_valid(self): + pass + + def test_get_hosts(self): + pass + + def test_is_group_valid(self): + pass + + def test_add_group(self): + pass + + def test_get_groups(self): + pass + + def test_get_group_relations(self): + pass + + def test_remove_group_relations(self): + pass + + def test_is_user_valid(self): + pass + + def test_get_users(self): + pass + + def test_get_user(self): + pass + + def test_get_domain_controllers(self): + pass + + def test_is_share_valid(self): + pass + + def test_add_share(self): + pass + + def test_get_shares(self): + pass + + def test_get_shares_by_access(self): + pass + + def test_get_users_with_share_access(self): + pass + + def test_add_domain_backupkey(self): + pass + + def test_get_domain_backupkey(self): + pass + + def test_is_dpapi_secret_valid(self): + pass + + def test_add_dpapi_secrets(self): + pass + + def test_get_dpapi_secrets(self): + pass + + def test_add_loggedin_relation(self): + pass + + def test_get_loggedin_relations(self): + pass + + def test_remove_loggedin_relations(self): + pass + From 2663b9e0d287dea8f8d140e669a8a6560ee85113 Mon Sep 17 00:00:00 2001 From: Marshall Hallenbeck Date: Mon, 20 Mar 2023 21:13:15 -0400 Subject: [PATCH 5/6] tests(deps): add pytest dependency --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index ca54abb9..c6d0ebe8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ asyauth = "^0.0.12" masky = "^0.2.0" sqlalchemy = "^2.0.4" aiosqlite = "^0.18.0" +pytest = "^7.2.2" [tool.poetry.dev-dependencies] flake8 = "*" From ef9700984ffa03985cdb2b87309d5b41d52258f5 Mon Sep 17 00:00:00 2001 From: Marshall Hallenbeck Date: Mon, 20 Mar 2023 21:14:07 -0400 Subject: [PATCH 6/6] update how workspaces are created so tests can utilize functionality --- cme/cmedb.py | 65 ++++++++++++++++++++++++++------------------- cme/crackmapexec.py | 2 +- cme/paths.py | 3 ++- 3 files changed, 40 insertions(+), 30 deletions(-) diff --git a/cme/cmedb.py b/cme/cmedb.py index b86dbbed..a5a5ab1a 100644 --- a/cme/cmedb.py +++ b/cme/cmedb.py @@ -3,6 +3,7 @@ import cmd import logging +import shutil import sqlite3 import sys import os @@ -10,7 +11,7 @@ import requests from terminaltables import AsciiTable import configparser from cme.loaders.protocol_loader import protocol_loader -from cme.paths import CONFIG_PATH, WS_PATH +from cme.paths import CONFIG_PATH, WS_PATH, WORKSPACE_DIR from requests import ConnectionError from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.exc import SAWarning @@ -303,7 +304,7 @@ class CMEDBMenu(cmd.Cmd): print("[-] Error reading cme.conf: {}".format(e)) sys.exit(1) - self.workspace_dir = os.path.expanduser('~/.cme/workspaces') + self.conn = None self.p_loader = protocol_loader() self.protocols = self.p_loader.get_protocols() @@ -323,7 +324,7 @@ class CMEDBMenu(cmd.Cmd): if not proto: return - proto_db_path = os.path.join(self.workspace_dir, self.workspace, proto + '.db') + proto_db_path = os.path.join(WORKSPACE_DIR, self.workspace, proto + '.db') if os.path.exists(proto_db_path): self.conn = create_db_engine(proto_db_path) db_nav_object = self.p_loader.load_protocol(self.protocols[proto]['nvpath']) @@ -354,38 +355,16 @@ class CMEDBMenu(cmd.Cmd): if subcommand == 'create': new_workspace = line.split()[1].strip() print("[*] Creating workspace '{}'".format(new_workspace)) - os.mkdir(os.path.join(self.workspace_dir, new_workspace)) - - for protocol in self.protocols.keys(): - try: - protocol_object = self.p_loader.load_protocol(self.protocols[protocol]['dbpath']) - except KeyError: - continue - proto_db_path = os.path.join(self.workspace_dir, new_workspace, protocol + '.db') - - if not os.path.exists(proto_db_path): - print('[*] Initializing {} protocol database'.format(protocol.upper())) - conn = sqlite3.connect(proto_db_path) - c = conn.cursor() - - # try to prevent some weird sqlite I/O errors - c.execute('PRAGMA journal_mode = OFF') - c.execute('PRAGMA foreign_keys = 1') - - getattr(protocol_object, 'database').db_schema(c) - - # commit the changes and close everything off - conn.commit() - conn.close() + self.create_workspace(new_workspace, self.p_loader, self.protocols) self.do_workspace(new_workspace) elif subcommand == 'list': print("[*] Enumerating Workspaces") - for workspace in os.listdir(os.path.join(self.workspace_dir)): + for workspace in os.listdir(os.path.join(WORKSPACE_DIR)): if workspace == self.workspace: print("==> "+workspace) else: print(workspace) - elif os.path.exists(os.path.join(self.workspace_dir, line)): + elif os.path.exists(os.path.join(WORKSPACE_DIR, line)): self.config.set('CME', 'workspace', line) self.write_configfile() self.workspace = line @@ -398,6 +377,36 @@ class CMEDBMenu(cmd.Cmd): print_help(help_string) +def create_workspace(workspace_name, p_loader, protocols): + os.mkdir(os.path.join(WORKSPACE_DIR, workspace_name)) + + for protocol in protocols.keys(): + try: + protocol_object = p_loader.load_protocol(protocols[protocol]['dbpath']) + except KeyError: + continue + proto_db_path = os.path.join(WORKSPACE_DIR, workspace_name, protocol + '.db') + + if not os.path.exists(proto_db_path): + print('[*] Initializing {} protocol database'.format(protocol.upper())) + conn = sqlite3.connect(proto_db_path) + c = conn.cursor() + + # try to prevent some weird sqlite I/O errors + c.execute('PRAGMA journal_mode = OFF') + c.execute('PRAGMA foreign_keys = 1') + + getattr(protocol_object, 'database').db_schema(c) + + # commit the changes and close everything off + conn.commit() + conn.close() + + +def delete_workspace(workspace_name): + shutil.rmtree(os.path.join(WORKSPACE_DIR, workspace_name)) + + def initialize_db(logger): if not os.path.exists(os.path.join(WS_PATH, 'default')): logger.info('Creating default workspace') diff --git a/cme/crackmapexec.py b/cme/crackmapexec.py index fd986c1b..10b71e48 100755 --- a/cme/crackmapexec.py +++ b/cme/crackmapexec.py @@ -228,7 +228,7 @@ def main(): 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}") + logging.debug(f"Protocol DB Object: {protocol_db_object}") db_path = os.path.join(CME_PATH, 'workspaces', current_workspace, args.protocol + '.db') logging.debug(f"DB Path: {db_path}") diff --git a/cme/paths.py b/cme/paths.py index 8c4a0eff..795bb088 100644 --- a/cme/paths.py +++ b/cme/paths.py @@ -9,4 +9,5 @@ if hasattr(sys, 'getandroidapilevel'): TMP_PATH = os.path.join('/data','data', 'com.termux', 'files', 'usr', 'tmp', 'cme_hosted') WS_PATH = os.path.join(CME_PATH, 'workspaces') CERT_PATH = os.path.join(CME_PATH, 'cme.pem') -CONFIG_PATH = os.path.join(CME_PATH, 'cme.conf') \ No newline at end of file +CONFIG_PATH = os.path.join(CME_PATH, 'cme.conf') +WORKSPACE_DIR = os.path.join(CME_PATH, 'workspaces')