diff --git a/docs/USAGE.md b/docs/USAGE.md index 753dcf59..7023692d 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -436,6 +436,6 @@ Every `Result` object returned by the engine supports the following methods: To format a **List** of results, use the `formatter`: - `into_json(results_list)` -- `into_csv(results_list)` +- `into_csv(results_list, include_header=True)` — pass `include_header=False` when appending to a CSV that already has a header, to avoid writing it twice. --- diff --git a/tests/test_formatter.py b/tests/test_formatter.py index 59f8017e..1d3d67f2 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -1,9 +1,10 @@ -from user_scanner.core.formatter import into_csv, into_json +from user_scanner.core.formatter import CSV_HEADER, into_csv, into_json from user_scanner.core.result import Result def test_get_result_output_formats(): res = Result.available(username="alice", site_name="ExampleSite", category="Cat") + res2 = Result.available(username="bob", site_name="ExampleSite", category="Cat") out_console = res.get_console_output() assert "Not Found" in out_console @@ -14,6 +15,23 @@ def test_get_result_output_formats(): assert '"username": "alice"' in out_json assert '"site_name": "ExampleSite"' in out_json + # Test CSV output with header included(Default behavior) out_csv = into_csv([res]) assert "alice" in out_csv assert "ExampleSite" in out_csv + + # Test CSV output with header excluded + out_csv_no_header = into_csv([res], include_header=False) + assert "alice" in out_csv_no_header + assert "ExampleSite" in out_csv_no_header + assert out_csv_no_header.count(CSV_HEADER) == 0 + + # Test CSV output with multiple results and header included + first = into_csv([res], include_header=True) + second = into_csv([res2], include_header=False) + + full = first + "\n" + second + + assert full.count(CSV_HEADER) == 1 + assert "alice" in full + assert "bob" in full diff --git a/user_scanner/__main__.py b/user_scanner/__main__.py index 3c2c188e..a14d2778 100644 --- a/user_scanner/__main__.py +++ b/user_scanner/__main__.py @@ -695,13 +695,14 @@ def main(): print(G + f"\n[+] JSON Results saved to {t_output}" + Style.RESET_ALL) elif args.format == "csv": - content_csv = formatter.into_csv(t_results) try: with open(t_output, "r", encoding="utf-8") as init_file: has_content = init_file.read().strip() != "" except Exception: has_content = False + content_csv = formatter.into_csv(t_results, include_header=not has_content) + with open(t_output, "a", encoding="utf-8") as f: if has_content: f.write("\n") diff --git a/user_scanner/core/formatter.py b/user_scanner/core/formatter.py index 28364e4b..5d0a73ef 100644 --- a/user_scanner/core/formatter.py +++ b/user_scanner/core/formatter.py @@ -15,8 +15,8 @@ def get_json_data(results: List[Result]) -> list: return [r.to_dict() for r in results] -def into_csv(results: List[Result]) -> str: - return CSV_HEADER + "\n" + "\n".join(result.to_csv() for result in results) +def into_csv(results: List[Result], include_header: bool=True) -> str: + return CSV_HEADER + "\n" + "\n".join(result.to_csv() for result in results) if include_header else "\n".join(result.to_csv() for result in results) def into_pdf(