def update_artifact_by_id(conn, artifact_id: int, artifact: Artifact) -> None:
"""
Update an artifact in the database
Arguments:
conn - An sqlite connection object representing
the database on which data will be inserted
artifact_id - the ID associated to the artifact to update
artifact - an artifact object which contains the attributes
to update
Returns:
Returns an error if there was a failure in the update operation
"""
current_artifact = get_artifact_by_id(conn, artifact_id)
if not current_artifact:
return None
update_record = (artifact_id, artifact.title, artifact.category,
artifact.path, artifact.tags, artifact.author,
artifact.status, None)
new_record = list()
for i, elem in enumerate(attr.astuple(current_artifact)):
new_record.append(update_record[i] or elem or None)
delete_artifact_by_id(conn, artifact_id)
updated_artifact = Artifact(None, *new_record[1:])
insert_artifact(conn, updated_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 get_artifacts_by_category(conn,
query_string: str = "",
is_strict: bool = False) -> List[Artifact]:
"""
Returns artifacts of the category matching the string to search,
if the string is empty, it will retrieve all artifacts
Arguments:
conn - the sqlite3 Connection object
query_string - the string to match in the database
is_strict - if True, the pattern matching is
on exact strings
Returns:
a set of the artifacts matching the provided string that
have been found
"""
if not is_strict:
query_string = "%" + query_string + "%"
sql_query = """SELECT *
FROM artifacts
WHERE category LIKE ?
COLLATE NOCASE"""
cur = conn.cursor()
cur.execute(sql_query, [query_string])
artifacts = [Artifact(*row) for row in cur.fetchall()]
return artifacts
def add_file_to_kb(conn, args: Dict[str, str], config: Dict[str, str],
fname: str) -> None:
"""
Adds a file to the kb knowledge base.
Arguments:
conn - the connection to the database object
args - the args dictionary passed to the add command,
it must contain at least the following keys:
title, category, tags, status, author
config - the configuration dictionary that must contain
at least the following key:
PATH_KB_DATA, the path to where artifact are stored
fname - the path of the file to add to kb
"""
title = args["title"] or fs.get_basename(fname)
category = args["category"] or "default"
category_path = Path(config["PATH_KB_DATA"], category)
category_path.mkdir(parents=True, exist_ok=True)
fs.copy_file(fname, Path(category_path, title))
if not db.is_artifact_existing(conn, title, category):
fs.copy_file(fname, Path(category_path, title))
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 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 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_delete_artifact_by_id():
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"))
db.delete_artifact_by_id(conn, 1)
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')}
db.delete_artifact_by_id(conn, 2)
sql = "SELECT * FROM artifacts;"
cur = conn.cursor()
cur.execute(sql)
rows = cur.fetchall()
assert len(rows) == 0
db_path.unlink()
def update_artifact(conn, old_artifact: Artifact, args: Dict[str, str], config: Dict[str, str], attachment):
"""
Update artifact properties within the knowledge base of kb.
Arguments:
old_artifact: - an object of type Artifact containing the old artifact details
args: - a dictionary containing the following fields:
id -> an id of an artifact - note - the ACTUAL db_id
title -> the title to be assigned to the artifact
to update
category -> the category to be assigned to the
artifact to update
tags -> the tags to be assigned to the artifact
to update
author -> the author to be assigned to the artifact
to update
status -> the status to be assigned to the artifact
to update
template -> the template to be assigned to the artifact
to update
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
attachment: - new file content
"""
initializer.init(config)
template_name = args.get("template", "")
updated_artifact = Artifact(
id=None,
title=args.get("title", old_artifact.title),
category=args.get("category", old_artifact.category),
tags=args.get("tags", old_artifact.tags),
author=args.get("author", old_artifact.author),
status=args.get("status", old_artifact.status),
template=args.get("template", old_artifact.template),
path=args.get("category", old_artifact.category) + '/' + args.get("title", old_artifact.title)
)
db.update_artifact_by_id(conn, old_artifact.id, updated_artifact)
# If either title or category has been changed, we must move the file
if args["category"] or args["title"]:
old_category_path = Path(
config["PATH_KB_DATA"],
old_artifact.category)
new_category_path = Path(
config["PATH_KB_DATA"],
args["category"] or old_artifact.category)
fs.create_directory(new_category_path)
fs.move_file(Path(old_category_path, old_artifact.title), Path(
new_category_path, args["title"] or old_artifact.title))
return -200
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 get_artifact_by_id(conn, artifact_id: int) -> Artifact:
"""
Get the artifact corresponding to the provided
artifact ID
Arguments:
conn - the sqlite3 connection object
artifact_id - the ID associated to the artifact to retrieve
"""
cur = conn.cursor()
sql_query = "SELECT * FROM artifacts WHERE id = ?"
cur.execute(sql_query, [artifact_id])
return Artifact(*cur.fetchall().pop())
def test_is_artifact_existing():
db_path = Path("tests", "data", "test_exist.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"))
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="")
def get_artifacts_by_tags(
conn,
tags: List[str] = [],
is_strict: bool = False
) -> List[Artifact]:
"""
Returns artifacts matching the provided list of tags,
if no tags are provided, it will retrieve all artifacts
Arguments:
conn - the sqlite3 Connection object
query_string - the list of tags (strings) provided
is_strict - if True, the pattern matching is
on exact strings
Returns:
a list of the found artifacts matching the provided tags
"""
rows = list()
if not is_strict:
tags = ["%" + tag + "%" for tag in tags]
cur = conn.cursor()
for tag in tags:
sql_query = """
SELECT *
FROM artifacts INNER JOIN tags ON tags.artifact_id = artifacts.id
WHERE tag LIKE ?
COLLATE NOCASE
"""
cur.execute(sql_query, [tag])
rows += cur.fetchall()
# We discard the last two fields in result rows
# since we are only interested in artifact data
rows = [r[:-2] for r in list(set(rows))]
artifacts = [Artifact(*row) for row in rows]
return artifacts
def update(args: Dict[str, str], config: Dict[str, str]):
"""
Update artifact properties within the knowledge base of kb.
Arguments:
args: - a dictionary containing the following fields:
id -> a list of IDs (the ones you see with kb list)
associated to the artifact to update
title -> the title to be assigned to the artifact
to update
category -> the category to be assigned to the
artifact to update
tags -> the tags to be assigned to the artifact
to update
author -> the author to be assigned to the artifact
to update
status -> the status to be assigned to the artifact
to update
template -> the template to be assigned to the artifact
to update
edit_content -> a boolean, if True -> also open the
artifact to edit the content
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"]:
old_artifact = history.get_artifact(conn, config["PATH_KB_HIST"],
args["id"])
if not old_artifact:
print("The artifact you are trying to update does not exist! "
"Please insert a valid ID...")
return None
updated_artifact = Artifact(id=None,
title=args["title"],
category=args["category"],
tags=args["tags"],
author=args["author"],
status=args["status"],
template=args["template"])
db.update_artifact_by_id(conn, old_artifact.id, updated_artifact)
# If either title or category has been changed, we must move the file
if args["category"] or args["title"]:
old_category_path = Path(config["PATH_KB_DATA"],
old_artifact.category)
new_category_path = Path(config["PATH_KB_DATA"], args["category"]
or old_artifact.category)
fs.create_directory(new_category_path)
fs.move_file(
Path(old_category_path, old_artifact.title),
Path(new_category_path, args["title"] or old_artifact.title))
# else if a title is specified
elif args["title"]:
artifact = db.get_uniq_artifact_by_filter(conn,
title=args["title"],
category=args["category"],
author=args["author"],
status=args["status"],
is_strict=True)
if artifact:
category_path = Path(config["PATH_KB_DATA"], artifact.category)
else:
print(
"There is none or more than one artifact with that title, please specify a category"
)
if args["edit_content"] or args["body"]:
if args["title"]:
artifact_path = str(Path(category_path, artifact.title))
shell_cmd = shlex.split(config["EDITOR"]) + [artifact_path]
elif args["id"]:
artifact_path = str(
Path(config["PATH_KB_DATA"]) / old_artifact.category /
old_artifact.title)
shell_cmd = shlex.split(config["EDITOR"]) + [artifact_path]
if args["body"]:
args["body"] = args["body"].replace("\\n", "\n")
with open(artifact_path, 'w') as art_file:
art_file.write(args["body"])
else:
call(shell_cmd)
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
