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 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 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_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 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 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)
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