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