def test_delete_artifact_by_name(): db_path = Path("tests","data","newdb.db") db.create_kb_database(db_path) conn = db.create_connection(db_path) db.insert_artifact(conn, path="pentest/smb", title="pentest_smb", category="procedure", tags=['pt','smb'], status="OK", author="gnc") db.insert_artifact(conn, path="protocol/ftp", title="ftp", category="cheatsheet", tags=[], status="Draft", author="elektroniz") db.delete_artifact_by_name(conn, title="pentest_smb", category="") sql = "SELECT * FROM artifacts;" cur = conn.cursor() cur.execute(sql) rows = cur.fetchall() assert len(rows) == 2 db.delete_artifact_by_name(conn, title="pentest_smb", category="procedure") sql = "SELECT * FROM artifacts;" cur = conn.cursor() cur.execute(sql) rows = cur.fetchall() assert len(rows) == 1 assert set(rows) == {(2, 'ftp', 'cheatsheet', 'protocol/ftp', '', 'Draft', 'elektroniz')} db_path.unlink()
def test_get_artifacts_by_tags(): db_path = Path("tests","data","kb_art_tags.db") conn = db.create_connection(db_path) db.create_kb_database(db_path) db.insert_artifact(conn, path="cheatsheet/pentest_smb", title="pentest_smb", category="procedure", tags=['pt','smb'], status="ok", author="gnc") db.insert_artifact(conn, path="guides/ftp", title="ftp", category="cheatsheet", tags=[], status="draft", author="elektroniz") db.insert_artifact(conn, path="guides/http", title="http", category="cheatsheet", status="OK", author="elektroniz") db.insert_artifact(conn, path="guides/irc", title="irc", category="cheatsheet", tags=["protocol"], status="draft", author="elektroniz") db.insert_artifact(conn, path="cheatsheet/pentest_ftp", title="pentest_ftp", category="cheatsheet", tags=["pt"], status="draft", author="elektroniz") rows = db.get_artifacts_by_tags(conn, tags=["pt"], is_strict=False) assert len(rows) == 2 rows = db.get_artifacts_by_tags(conn, tags=["p"], is_strict=False) assert len(rows) == 3 rows = db.get_artifacts_by_tags(conn, tags=["pt"], is_strict=True) assert len(rows) == 2 db_path.unlink()
def test_is_artifact_existing(): db_path = Path("tests", "data", "newdb.db") db.create_kb_database(db_path) conn = db.create_connection(db_path) db.insert_artifact( conn, Artifact(id=None, path="pentest/smb", title="pentest_smb", category="procedure", tags='pt;smb', status="OK", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="protocol/ftp", title="ftp", category="cheatsheet", status="Draft", author="elektroniz")) assert db.is_artifact_existing(conn, title="pentest_smb", category="procedure") assert db.is_artifact_existing(conn, title="ftp", category="cheatsheet") assert not db.is_artifact_existing( conn, title="pentest_smb", category="nonexist") assert not db.is_artifact_existing( conn, title="nonexist", category="procedure") assert not db.is_artifact_existing(conn, title="", category="cheatsheet") assert not db.is_artifact_existing(conn, title="", category="") db_path.unlink()
def test_get_artifacts_by_filter(): db_path = Path("tests","data","kb_filter.db") conn = db.create_connection(db_path) db.create_kb_database(db_path) db.insert_artifact(conn, path="", title="pentest_smb", category="procedure", tags=['pt','smb'], status="ok", author="gnc") db.insert_artifact(conn, path="", title="ftp", category="cheatsheet", tags=["protocol"], status="draft", author="elektroniz") db.insert_artifact(conn, path="", title="pentest_ftp", category="procedure", tags=["pt","ftp"], status="draft", author="elektroniz") db.insert_artifact(conn, path="general/CORS", title="CORS", category="general", tags=["web"], status="draft", author="elektroniz") rows = db.get_artifacts_by_filter(conn, title="pentest", category="cheatsheet", tags=["pt"], is_strict=False) assert len(rows) == 0 rows = db.get_artifacts_by_filter(conn, category="procedure", tags=["pt"], is_strict=False) print(rows) assert set(rows) == {(1,"pentest_smb","procedure","","pt;smb","ok","gnc"), (3,"pentest_ftp","procedure","","pt;ftp","draft", "elektroniz")} rows = db.get_artifacts_by_filter(conn, title="OR") assert set(rows) == {(4,"CORS","general","general/CORS","web", "draft","elektroniz")} rows = db.get_artifacts_by_filter(conn, category="cheatsheet", is_strict=False) assert set(rows) == {(2,"ftp","cheatsheet","","protocol", "draft", "elektroniz")} rows = db.get_artifacts_by_filter(conn, category="sheet", is_strict=False) assert set(rows) == {(2,"ftp","cheatsheet","","protocol", "draft", "elektroniz")} rows = db.get_artifacts_by_filter(conn, category="cheatsheet", is_strict=True) assert set(rows) == {(2,"ftp","cheatsheet","","protocol", "draft", "elektroniz")} rows = db.get_artifacts_by_filter(conn, category="sheet", is_strict=True) assert len(rows) == 0 db_path.unlink()
def test_get_artifacts_by_tags(): db_path = Path("tests", "data", "kb_art_tags.db") conn = db.create_connection(str(db_path)) with conn: schema_version = 1 db.create_kb_database(str(db_path), schema_version) db.insert_artifact( conn, Artifact(id=None, path="cheatsheet/pentest_smb", title="pentest_smb", category="procedure", tags='pt;smb', status="ok", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="guides/ftp", title="ftp", category="cheatsheet", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="guides/http", title="http", category="cheatsheet", status="OK", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="guides/irc", title="irc", category="cheatsheet", tags="protocol", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="cheatsheet/pentest_ftp", title="pentest_ftp", category="cheatsheet", tags="pt", status="draft", author="elektroniz")) rows = db.get_artifacts_by_tags(conn, tags=["pt"], is_strict=False) assert len(rows) == 2 rows = db.get_artifacts_by_tags(conn, tags=["p"], is_strict=False) assert len(rows) == 3 rows = db.get_artifacts_by_tags(conn, tags=["pt"], is_strict=True) assert len(rows) == 2
def create_kb_files(config): """ Create the kb files and infrastructure Arguments: config - a dictionary containing the following keys: PATH_KB - the path to kb (~/.kb by default) PATH_KB_DB - the path to kb database (~/.kb/kb.db by default) PATH_KB_DATA - the path to kb data (~/.kb/data/ by default) INITIAL_CATEGORIES - a list containing the initial categories contained within kb PATH_KB_TEMPLATES - the path to kb templates (~/.kb/templates/ by default) DB_SCHEMA_VERSION - the database schema version """ # Get paths for kb from configuration kb_path = config["PATH_KB"] db_path = config["PATH_KB_DB"] data_path = config["PATH_KB_DATA"] initial_categs = config["INITIAL_CATEGORIES"] templates_path = config["PATH_KB_TEMPLATES"] schema_version = config["DB_SCHEMA_VERSION"] default_template_path = str(Path(templates_path) / "default") # Create main kb fs.create_directory(kb_path) # Create kb database if not os.path.exists(db_path): db.create_kb_database(db_path, schema_version) # Check schema version conn = db.create_connection(db_path) current_schema_version = db.get_schema_version(conn) if current_schema_version == 0: db.migrate_v0_to_v1(conn) # Create "data" directory fs.create_directory(data_path) # Create "templates" directory fs.create_directory(templates_path) # Create kb initial categories directories for category in initial_categs: category_path = Path(data_path, category) fs.create_directory(category_path) # Create markers file with open(default_template_path, 'w') as cfg: cfg.write(toml.dumps(conf.DEFAULT_TEMPLATE))
def test_get_artifacts_by_title(): db_path = Path("tests", "data", "kb_filter_title.db") conn = db.create_connection(db_path) db.create_kb_database(db_path) db.insert_artifact( conn, Artifact(id=None, path="cheatsheet/pentest_smb", title="pentest_smb", category="procedure", tags='pt;smb', status="ok", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="guides/ftp", title="ftp", category="cheatsheet", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="guides/http", title="http", category="cheatsheet", status="OK", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="guides/irc", title="irc", category="cheatsheet", tags="protocol", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="cheatsheet/pentest_ftp", title="pentest_ftp", category="cheatsheet", tags="pt", status="draft", author="elektroniz")) rows = db.get_artifacts_by_title(conn, query_string="", is_strict=False) assert len(rows) == 5 rows = db.get_artifacts_by_title(conn, query_string="", is_strict=True) assert len(rows) == 0 db_path.unlink()
def test_create_kb_database_table(): db_path = Path("tests","data","mydb.db") db.create_kb_database(db_path) conn = db.create_connection(db_path) try: c = conn.cursor() except Exception as e: print(e) # Make sure, the database has the correct number of tables kb_tables = _list_tables(conn) assert len(kb_tables) == 2 assert kb_tables == [("artifacts",),("tags",)] db_path.unlink()
def test_create_kb_database_table(): db_path = Path("tests", "data", "mydb3.db") schema_version = 1 db.create_kb_database(str(db_path), schema_version) conn = db.create_connection(str(db_path)) with conn: try: c = conn.cursor() except Exception as e: print(e) # Make sure, the database has the correct number of tables kb_tables = _list_tables(conn) assert len(kb_tables) == 2 assert kb_tables == [("artifacts", ), ("tags", )]
def test_delete_artifact_by_name(): db_path = Path("tests", "data", "test_name.db") schema_version = 1 db.create_kb_database(str(db_path), schema_version) conn = db.create_connection(str(db_path)) with conn: db.insert_artifact( conn, Artifact(id=None, path="pentest/smb", title="pentest_smb", category="procedure", tags='pt;smb', status="OK", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="protocol/ftp", title="ftp", category="cheatsheet", status="Draft", author="elektroniz")) db.delete_artifact_by_name(conn, title="pentest_smb", category="") sql = "SELECT * FROM artifacts;" cur = conn.cursor() cur.execute(sql) rows = cur.fetchall() assert len(rows) == 2 db.delete_artifact_by_name(conn, title="pentest_smb", category="procedure") sql = "SELECT * FROM artifacts;" cur = conn.cursor() cur.execute(sql) rows = cur.fetchall() assert len(rows) == 1 assert set(rows) == {(2, 'ftp', 'cheatsheet', 'protocol/ftp', None, 'Draft', 'elektroniz', None)}
def test_get_artifacts_by_category(): db_path = Path("tests","data","kb_filter_cat.db") conn = db.create_connection(db_path) db.create_kb_database(db_path) db.insert_artifact(conn, path="cheatsheet/pentest_smb", title="pentest_smb", category="procedure", tags=['pt','smb'], status="ok", author="gnc") db.insert_artifact(conn, path="guides/ftp", title="ftp", category="cheatsheet", tags=[], status="draft", author="elektroniz") db.insert_artifact(conn, path="guides/http", title="http", category="cheatsheet", status="OK", author="elektroniz") db.insert_artifact(conn, path="guides/irc", title="irc", category="cheatsheet", tags=["protocol"], status="draft", author="elektroniz") db.insert_artifact(conn, path="cheatsheet/pentest_ftp", title="pentest_ftp", category="cheatsheet", tags=["pt"], status="draft", author="elektroniz") db.insert_artifact(conn, path="sheet/math", title="math_formulas", category="sheet", tags=["math"], status="draft", author="gnc") db.insert_artifact(conn, path="sheet/math2", title="geometry_formulas", category="sheet", tags=["math"], status="draft", author="gnc") rows = db.get_artifacts_by_category(conn, query_string="", is_strict=False) assert len(rows) == 7 rows = db.get_artifacts_by_category(conn, query_string="", is_strict=True) assert len(rows) == 0 rows = db.get_artifacts_by_category(conn, query_string="sheet", is_strict=True) assert len(rows) == 2 db_path.unlink()
def create_kb_files(config): """ Create the kb files and infrastructure Arguments: config - a dictionary containing the following keys: PATH_KB - the path to kb (~/.kb by default) PATH_KB_DB - the path to kb database (~/.kb/kb.db by default) PATH_KB_DATA - the path to kb data (~/.kb/data/ by default) PATH_KB_MARKERS - the path to kb markers (~/.kb/markers.toml by default) """ # Get paths for kb from configuration kb_path = config["PATH_KB"] db_path = config["PATH_KB_DB"] data_path = config["PATH_KB_DATA"] markers_path = config["PATH_KB_MARKERS"] initial_categs = config["INITIAL_CATEGORIES"] # Create main kb fs.create_directory(kb_path) # Create kb database if not os.path.exists(db_path): db.create_kb_database(db_path) # Create "data" director fs.create_directory(data_path) # Create kb initial categories directories for category in initial_categs: category_path = Path(data_path, category) fs.create_directory(category_path) # Create markers file with open(markers_path, 'w') as cfg: cfg.write(toml.dumps(conf.DEFAULT_MARKERS))
def test_insert_artifact(): db_path = Path("tests", "data", "test_insert.db") schema_version = 1 db.create_kb_database(str(db_path), schema_version) conn = db.create_connection(str(db_path)) with conn: db.insert_artifact( conn, Artifact(id=None, path="pentest/smb", title="pentest_smb", category="procedure", tags='pt;smb', status="OK", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="protocol/ftp", title="ftp", category="cheatsheet", status="Draft", author="elektroniz")) kb_tables = _list_tables(conn) assert len(kb_tables) == 2 assert kb_tables == [("artifacts", ), ("tags", )] sql = "SELECT * FROM artifacts;" cur = conn.cursor() cur.execute(sql) rows = cur.fetchall() print(rows) assert set(rows) == {(1, 'pentest_smb', 'procedure', 'pentest/smb', 'pt;smb', 'OK', 'gnc', None), (2, 'ftp', 'cheatsheet', 'protocol/ftp', None, 'Draft', 'elektroniz', None)}
def test_insert_artifact(): db_path = Path("tests","data","newdb.db") db.create_kb_database(db_path) conn = db.create_connection(db_path) db.insert_artifact(conn, path="pentest/smb", title="pentest_smb", category="procedure", tags=['pt','smb'], status="OK", author="gnc") db.insert_artifact(conn, path="protocol/ftp", title="ftp", category="cheatsheet", tags=[], status="Draft", author="elektroniz") kb_tables = _list_tables(conn) assert len(kb_tables) == 2 assert kb_tables == [("artifacts",),("tags",)] sql = "SELECT * FROM artifacts;" cur = conn.cursor() cur.execute(sql) rows = cur.fetchall() assert set(rows) == {(1, 'pentest_smb', 'procedure', 'pentest/smb', 'pt;smb', 'OK', 'gnc'), (2, 'ftp', 'cheatsheet', 'protocol/ftp', '', 'Draft', 'elektroniz')} db_path.unlink()
def test_get_artifacts_by_filter(): db_path = Path("tests", "data", "kb_filter.db") conn = db.create_connection(str(db_path)) with conn: schema_version = 1 db.create_kb_database(str(db_path), schema_version) db.insert_artifact( conn, Artifact(id=None, path="", title="pentest_smb", category="procedure", tags='pt;smb', status="ok", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="", title="ftp", category="cheatsheet", tags="protocol", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="", title="pentest_ftp", category="procedure", tags="pt;ftp", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="general/CORS", title="CORS", category="general", tags="web", status="draft", author="elektroniz")) rows = db.get_artifacts_by_filter(conn, title="pentest", category="cheatsheet", tags=["pt"], is_strict=False) assert len(rows) == 0 rows = db.get_artifacts_by_filter(conn, category="procedure", tags=["pt"], is_strict=False) assert sorted(list(set(rows)), key=lambda i: i.id) == [ Artifact(1, "pentest_smb", "procedure", "procedure/pentest_smb", "pt;smb", "ok", "gnc", None), Artifact(3, "pentest_ftp", "procedure", "procedure/pentest_ftp", "pt;ftp", "draft", "elektroniz", None) ] rows = db.get_artifacts_by_filter(conn, title="OR") assert set(rows) == { Artifact(4, "CORS", "general", "general/CORS", "web", "draft", "elektroniz", None) } rows = db.get_artifacts_by_filter(conn, category="cheatsheet", is_strict=False) assert set(rows) == { Artifact(2, "ftp", "cheatsheet", "cheatsheet/ftp", "protocol", "draft", "elektroniz", None) } rows = db.get_artifacts_by_filter(conn, category="sheet", is_strict=False) assert set(rows) == { Artifact(2, "ftp", "cheatsheet", "cheatsheet/ftp", "protocol", "draft", "elektroniz", None) } rows = db.get_artifacts_by_filter(conn, category="cheatsheet", is_strict=True) assert set(rows) == { Artifact(2, "ftp", "cheatsheet", "cheatsheet/ftp", "protocol", "draft", "elektroniz", None) } rows = db.get_artifacts_by_filter(conn, category="sheet", is_strict=True) assert len(rows) == 0
def test_get_artifacts_by_category(): db_path = Path("tests", "data", "kb_filter_cat.db") conn = db.create_connection(str(db_path)) with conn: schema_version = 1 db.create_kb_database(str(db_path), schema_version) db.insert_artifact( conn, Artifact(id=None, path="cheatsheet/pentest_smb", title="pentest_smb", category="procedure", tags='pt;smb', status="ok", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="guides/ftp", title="ftp", category="cheatsheet", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="guides/http", title="http", category="cheatsheet", status="OK", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="guides/irc", title="irc", category="cheatsheet", tags="protocol", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="cheatsheet/pentest_ftp", title="pentest_ftp", category="cheatsheet", tags="pt", status="draft", author="elektroniz")) db.insert_artifact( conn, Artifact(id=None, path="sheet/math", title="math_formulas", category="sheet", tags="math", status="draft", author="gnc")) db.insert_artifact( conn, Artifact(id=None, path="sheet/math2", title="geometry_formulas", category="sheet", tags="math", status="draft", author="gnc")) rows = db.get_artifacts_by_category(conn, query_string="", is_strict=False) assert len(rows) == 7 rows = db.get_artifacts_by_category(conn, query_string="", is_strict=True) assert len(rows) == 0 rows = db.get_artifacts_by_category(conn, query_string="sheet", is_strict=True) assert len(rows) == 2