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 edit_by_name(title: str, category: str, config: Dict[str, str]):
"""
Edit the content of an artifact by name, that is title/category
Arguments:
title: - the title assigned to the artifact(s)
category: - the category assigned to the artifact(s)
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
"""
conn = db.create_connection(config["PATH_KB_DB"])
artifacts = db.get_artifacts_by_filter(conn, title=title,
category=category,
is_strict=True)
if len(artifacts) == 1:
artifact = artifacts.pop()
category_path = Path(config["PATH_KB_DATA"], artifact.category)
shell_cmd = shlex.split(
config["EDITOR"]) + [str(Path(category_path, artifact.title))]
call(shell_cmd)
elif len(artifacts) > 1:
print(
"There is more than one artifact with that title, please specify a category")
else:
print(
"There is no artifact with that name, please specify a correct artifact name")
def is_initialized(config) -> bool:
"""
Check if kb is correctly initialized,
ensure that:
1 - the .kb directory exists
2 - the kb database exists
3 - the kb data directory exists
Arguments:
config - a dictionary containing the following keys:
PATH_KB - the path to kb
(~/.kb by default)
PATH_KB_DB - the path to kb
(~/.kb/kb.db by default)
PATH_KB_DATA - the path to kb
(~/.kb/data/ by default)
Returns:
True is kb is correctly initialized, False otherwise
"""
kb_path = config["PATH_KB"]
db_path = config["PATH_KB_DB"]
data_path = config["PATH_KB_DATA"]
templates_path = config["PATH_KB_TEMPLATES"]
for path in [kb_path, db_path, data_path, templates_path]:
if not os.path.exists(path):
return False
conn = db.create_connection(db_path)
return db.is_schema_updated_to_version(conn, config["DB_SCHEMA_VERSION"])
def add(args: Dict[str, str], config: Dict[str, str]):
"""
Adds a list of artifacts to the knowledge base of kb.
Arguments:
args: - a dictionary containing the following fields:
file -> a list of files to add to kb
title -> the title assigned to the artifact(s)
category -> the category assigned to the artifact(s)
tags -> the tags assigned to the artifact(s)
author -> the author to assign to the artifact
status -> the status to assign to the artifact
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
EDITOR - the editor program to call
"""
# Check if the add command has proper arguments/options
is_valid_add = args["file"] or args["title"]
if not is_valid_add:
print("Please, either specify a file or a title for the new artifact")
sys.exit(1)
# Check initialization
initializer.init(config)
conn = db.create_connection(config["PATH_KB_DB"])
if args["file"]:
for fname in args["file"]:
if fs.is_directory(fname):
continue
add_file_to_kb(conn, args, config, fname)
else:
# Get title for the new artifact
title = args["title"]
# Assign a "default" category if not provided
category = args["category"] or "default"
# Create "category" directory if it does not exist
category_path = Path(config["PATH_KB_DATA"], category)
category_path.mkdir(parents=True, exist_ok=True)
if not db.is_artifact_existing(conn, title, category):
# If a file is provided, copy the file to kb directory
# otherwise open up the editor and create some content
shell_cmd = shlex.split(
config["EDITOR"]) + [str(Path(category_path, title))]
call(shell_cmd)
new_artifact = Artifact(id=None,
title=title,
category=category,
path="{category}/{title}".format(
category=category, title=title),
tags=args["tags"],
status=args["status"],
author=args["author"])
db.insert_artifact(conn, new_artifact)
def apply_on_set(args: Dict[str, str], config: Dict[str, str]):
"""
Apply the specified template to all the filtered artifacts
"""
# Check initialization
initializer.init(config)
tags_list = None
if args["tags"] and args["tags"] != "":
tags_list = args["tags"].split(';')
conn = db.create_connection(config["PATH_KB_DB"])
is_query_strict = not args["extended_match"]
rows = db.get_artifacts_by_filter(conn,
title=args["title"],
category=args["category"],
tags=tags_list,
status=args["status"],
author=args["author"],
is_strict=is_query_strict)
for artifact in rows:
updated_artifact = Artifact(id=artifact.id,
title=artifact.title,
category=artifact.category,
tags=artifact.tags,
author=artifact.author,
status=artifact.status,
template=args["template"])
db.update_artifact_by_id(conn, artifact.id, updated_artifact)
def delete_by_id(id: int, config: Dict[str, str]):
"""
Edit the content of an artifact by id.
Arguments:
id: - the ID (the one you see with kb list)
associated to the artifact to delete
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
"""
conn = db.create_connection(config["PATH_KB_DB"])
artifact_id = history.get_artifact_id(config["PATH_KB_HIST"], id)
artifact = db.get_artifact_by_id(conn, artifact_id)
if not artifact:
return
db.delete_artifact_by_id(conn, artifact_id)
category_path = Path(config["PATH_KB_DATA"], artifact.category)
try:
Path(category_path, artifact.title).unlink()
except FileNotFoundError:
pass
if fs.count_files(category_path) == 0:
fs.remove_directory(category_path)
print("Artifact {category}/{title} removed!".format(
category=artifact.category, title=artifact.title))
def delete_by_name(title: str, category: str, config: Dict[str, str]):
"""
Edit the content of an artifact by name, that is title/category
Arguments:
title: - the title assigned to the artifact to delete
category: - the category assigned to the artifact to delete
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
"""
conn = db.create_connection(config["PATH_KB_DB"])
artifacts = db.get_artifacts_by_filter(conn, title=title,
category=category,
is_strict=True)
if len(artifacts) == 1:
artifact = artifacts.pop()
db.delete_artifact_by_id(conn, artifact.id)
print("Artifact {}/{} removed!".format(artifact.category, artifact.title))
elif len(artifacts) > 1:
print(
"There is more than one artifact with that title, please specify a category")
else:
print(
"There is no artifact with that name, please specify a correct artifact name")
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_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_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 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 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 edit(args: Dict[str, str], config: Dict[str, str]):
"""
Edit the content of an artifact.
Arguments:
args: - a dictionary containing the following fields:
id -> the IDs (the one you see with kb list)
associated to the artifact we want to edit
title -> the title assigned to the artifact(s)
category -> the category assigned to the artifact(s)
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
"""
initializer.init(config)
conn = db.create_connection(config["PATH_KB_DB"])
# if an ID is specified, load artifact with that ID
if args["id"]:
artifact = history.get_artifact(conn, config["PATH_KB_HIST"],
args["id"])
category_path = Path(config["PATH_KB_DATA"], artifact.category)
shell_cmd = config["EDITOR"].split() + [
Path(category_path, artifact.title)
]
call(shell_cmd)
# else if a title is specified
elif args["title"]:
artifacts = db.get_artifacts_by_filter(conn,
title=args["title"],
category=args["category"],
is_strict=True)
if len(artifacts) == 1:
artifact = artifacts.pop()
category_path = Path(config["PATH_KB_DATA"], artifact.category)
shell_cmd = config["EDITOR"].split() + [
Path(category_path, artifact.title)
]
call(shell_cmd)
elif len(artifacts) > 1:
print(
"There is more than one artifact with that title, please specify a category"
)
else:
print(
"There is no artifact with that name, please specify a correct artifact name"
)
def search_kb(args: Dict[str, str], config: Dict[str, str]):
"""
Search artifacts within the knowledge base of kb.
Arguments:
args: - a dictionary containing the following fields:
query -> filter for the title field of the artifact
category -> filter for the category field of the artifact
tags -> filter for the tags field of the artifact
author -> filter for the author field of the artifact
status -> filter for the status field of the artifact
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
"""
# Check initialization
initializer.init(config)
conn = db.create_connection(config["PATH_KB_DB"])
categories = ls.list_categories(config)
print(categories)
# List all categories
if args.get("all_categories", False) is True:
categories = ls.list_categories(config)
return categories
# List all categories
if args.get("all_tags", False) is True:
all_tags = ls.list_tags(conn, config)
return all_tags
tags_list = None
if args.get("tags",False) is True:
if args["tags"] and args["tags"] != "":
tags_list = args["tags"].split(';')
rows = db.get_artifacts_by_filter(
conn,
title=args.get("query",''),
category=args.get("category",''),
tags=tags_list,
status=args.get("status",''),
author=args.get("author",''))
artifacts = sorted(rows, key=lambda x: x.title)
return artifacts
def delete(args: Dict[str, str], config: Dict[str, str]):
"""
Delete a list of artifacts from the kb knowledge base.
Arguments:
args: - a dictionary containing the following fields:
id -> a list of IDs (the ones you see with kb list)
associated to the artifacts we want to delete
title -> the title assigned to the artifact(s)
category -> the category assigned to the artifact(s)
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
"""
initializer.init(config)
conn = db.create_connection(config["PATH_KB_DB"])
if args["id"]:
for i in args["id"]:
artifact_id = history.get_artifact_id(config["PATH_KB_HIST"], i)
artifact = db.get_artifact_by_id(conn, artifact_id)
if not artifact:
continue
db.delete_artifact_by_id(conn, artifact_id)
category_path = Path(config["PATH_KB_DATA"], artifact.category)
Path(category_path, artifact.title).unlink()
if fs.count_files(category_path) == 0:
fs.remove_directory(category_path)
print("Artifact {category}/{title} removed!".format(
category=artifact.category, title=artifact.title))
sys.exit(0)
# else if a title is specified
elif args["title"]:
artifacts = db.get_artifacts_by_filter(conn, title=args["title"],
category=args["category"],
is_strict=True)
if len(artifacts) == 1:
artifact = artifacts.pop()
db.delete_artifact_by_id(conn, artifact.id)
print("Artifact {}/{} removed!".format(artifact.category, artifact.title))
else:
print(
"There is more than one artifact with that title, please specify a category")
def view_by_name(title: str, category: str, config: Dict[str, str],
open_editor: bool, color_mode: bool):
"""
View the content of an artifact by name, that is title/category
Arguments:
title: - the title assigned to the artifact(s)
category: - the category assigned to the artifact(s)
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
open_editor - a boolean, if True it will open the artifact as
a temporary copy in editor
color_mode - a boolean, if True the colors on screen will be
enabled when printed on stdout
"""
conn = db.create_connection(config["PATH_KB_DB"])
artifacts = db.get_artifacts_by_filter(conn,
title=title,
category=category,
is_strict=True)
if len(artifacts) == 1:
artifact = artifacts.pop()
category_path = Path(config["PATH_KB_DATA"], artifact.category)
artifact_path = Path(category_path, artifact.title)
if open_editor:
tmpfname = fs.get_temp_filepath()
fs.copy_file(artifact_path, tmpfname)
shell_cmd = shlex.split(config["EDITOR"]) + [tmpfname]
call(shell_cmd)
fs.remove_file(tmpfname)
sys.exit(0)
# View File
if fs.is_text_file(artifact_path):
markers = get_template(artifact, config)
viewer.view(artifact_path, markers, color=color_mode)
else:
opener.open_non_text_file(artifact_path)
elif len(artifacts) > 1:
print(
"There is more than one artifact with that title, please specify a category"
)
else:
print(
"There is no artifact with that name, please specify a correct artifact name"
)
def search(args: Dict[str, str], config: Dict[str, str]):
"""
Search artifacts within the knowledge base of kb.
Arguments:
args: - a dictionary containing the following fields:
query -> filter for the title field of the artifact
category -> filter for the category field of the artifact
tags -> filter for the tags field of the artifact
author -> filter for the author field of the artifact
status -> filter for the status field of the artifact
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
"""
# Check initialization
initializer.init(config)
tags_list = None
if args["tags"] and args["tags"] != "":
tags_list = args["tags"].split(';')
conn = db.create_connection(config["PATH_KB_DB"])
rows = db.get_artifacts_by_filter(
conn,
title=args["query"],
category=args["category"],
tags=tags_list,
status=args["status"],
author=args["author"])
# rows.sort(key=lambda x: x[1])
artifacts = sorted(rows, key=lambda x: x.title)
# Write to history file
history.write(config["PATH_KB_HIST"], artifacts)
# Is full_identifier mode enabled?
if args["full_identifier"]:
printer.print_search_result_full_mode(artifacts)
return
# Print resulting list
color_mode = not args["no_color"]
if args["verbose"]:
printer.print_search_result_verbose(artifacts, color_mode)
else:
printer.print_search_result(artifacts, color_mode)
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 add(args: Dict[str, str], config: Dict[str, str]):
"""
Adds a list of artifacts to the knowledge base of kb.
Arguments:
args: - a dictionary containing the following fields:
file -> a list of files to add to kb
title -> the title assigned to the artifact(s)
category -> the category assigned to the artifact(s)
tags -> the tags assigned to the artifact(s)
author -> the author to assign to the artifact
status -> the status to assign to the artifact
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
EDITOR - the editor program to call
"""
# Check if the add command has proper arguments/options
is_valid_add = args["file"] or args["title"]
if not is_valid_add:
print("Please, either specify a file or a title for the new artifact")
sys.exit(1)
# Check initialization
initializer.init(config)
conn = db.create_connection(config["PATH_KB_DB"])
if args["file"]:
for fname in args["file"]:
if fs.is_directory(fname):
continue
add_file_to_kb(conn, args, config, fname)
else:
if not db.is_artifact_existing(conn, args["title"], args["category"]):
pass
else:
with tempfile.NamedTemporaryFile(delete=True) as f:
shell_cmd = shlex.split(config["EDITOR"]) + [f]
call(shell_cmd)
args["temp_file"] = f
result = add_artifact(conn, args, config)
return (result)
def view_by_id(id: int, config: Dict[str, str], open_editor: bool,
color_mode: bool):
"""
View the content of an artifact by id.
Arguments:
id: - the ID (the one you see with kb list)
associated to the artifact we want to edit
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
open_editor - a boolean, if True it will open the artifact as
a temporary copy in editor
color_mode - a boolean, if True the colors on screen will be
enabled when printed on stdout
"""
conn = db.create_connection(config["PATH_KB_DB"])
artifact_id = history.get_artifact_id(config["PATH_KB_HIST"], id)
artifact = db.get_artifact_by_id(conn, artifact_id)
if not artifact:
sys.exit(1)
category_path = Path(config["PATH_KB_DATA"], artifact.category)
artifact_path = Path(category_path, artifact.title)
if open_editor:
tmpfname = fs.get_temp_filepath()
fs.copy_file(artifact_path, tmpfname)
shell_cmd = shlex.split(config["EDITOR"]) + [tmpfname]
call(shell_cmd)
fs.remove_file(tmpfname)
sys.exit(0)
# View File
if fs.is_text_file(artifact_path):
markers = get_template(artifact, config)
viewer.view(artifact_path, markers, color=color_mode)
else:
opener.open_non_text_file(artifact_path)
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 test_create_table(db_connect):
sql_db_create_query = """CREATE TABLE IF NOT EXISTS artifacts (
id integer PRIMARY KEY,
title text NOT NULL,
category text NOT NULL,
path text NOT NULL,
tags text,
status text,
author text);
"""
db_path = Path("tests","data","mydb.db")
conn = db.create_connection(db_path)
if conn is not None:
db.create_table(conn, sql_db_create_query)
else:
print("Error! cannot create the database connection.")
assert len(_list_tables(conn)) == 1
db_path.unlink()
def test_list_tables():
sql_create_table_query = """CREATE TABLE IF NOT EXISTS testtable1 (
id integer PRIMARY KEY,
author text
);
CREATE TABLE IF NOT EXISTS testtable2 (
id integer PRIMARY KEY,
author text
);
"""
db_path = Path("tests","data","two_tables.db")
conn = db.create_connection(db_path)
try:
c = conn.cursor()
c.executescript(sql_create_table_query)
except Exception as e:
print(e)
assert len(_list_tables(conn)) == 2
db_path.unlink()
def edit_by_id(id: int, config: Dict[str, str]):
"""
Edit the content of an artifact by id.
Arguments:
id: - the ID (the one you see with kb list)
associated to the artifact to edit
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
EDITOR - the editor program to call
"""
conn = db.create_connection(config["PATH_KB_DB"])
artifact = history.get_artifact(conn, config["PATH_KB_HIST"], id)
category_path = Path(config["PATH_KB_DATA"], artifact.category)
shell_cmd = shlex.split(config["EDITOR"]) + \
[str(Path(category_path, artifact.title))]
call(shell_cmd)
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 grep(args: Dict[str, str], config: Dict[str, str]):
"""
Grep through the list of artifacts of the knowledge base of kb.
Arguments:
args: - a dictionary containing the following fields:
regex -> the regex to search for
case_insensitive -> a boolean, if true,
the search will be case insensitive
matches -> a boolean, if true, only the raw
matches will be shown
verbose -> a boolean, if true, a verbose
output is produced on screen
config: - a configuration dictionary containing at least
the following keys:
PATH_KB_DB - the database path of KB
PATH_KB_DATA - the data directory of KB
PATH_KB_HIST - the history menu path of KB
"""
initializer.init(config)
conn = db.create_connection(config["PATH_KB_DB"])
# Get all artifacts
rows = db.get_artifacts_by_filter(conn, title="")
# Get all the file paths related to the artifacts in the database
file_list = [Path(config["PATH_KB_DATA"], r.category, r.title)
for r in rows]
# Grep in the files
results = fs.grep_in_files(
file_list,
args["regex"],
args["case_insensitive"])
# Get the list of artifact tuples in the form (category,title)
artifact_names = [fs.get_filename_parts_wo_prefix(
res[0], config["PATH_KB_DATA"]) for res in results]
# If user specied --matches -> just show matching lines and exit
if args["matches"]:
printer.print_grep_matches(artifact_names)
sys.exit(0)
# Get the set of uniq artifacts
uniq_artifact_names = set(artifact_names)
# Get the number of matches (hits) for each path found
filecounts = get_hits_per_artifact_name(artifact_names)
grep_result = list()
for art in uniq_artifact_names:
artifact = db.get_artifacts_by_filter(
conn, category=art[0], title=art[1])[0]
if artifact:
no_of_hits = filecounts[art]
grep_result.append((artifact, no_of_hits))
# Sort by number of hits, the largest -> the first
grep_result.sort(key=lambda x: x[1], reverse=True)
grep_artifacts = [r[0] for r in grep_result]
grep_hits = [r[1] for r in grep_result]
# Write to history file
history.write(config["PATH_KB_HIST"], grep_artifacts)
color_mode = not args["no_color"]
if args["verbose"]:
printer.print_grep_result_verbose(
grep_artifacts, grep_hits, color_mode)
else:
printer.print_grep_result(grep_artifacts, grep_hits, color_mode)
