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
