diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py index e954261..d03429a 100644 --- a/parsedmarc/__init__.py +++ b/parsedmarc/__init__.py @@ -1255,7 +1255,11 @@ def watch_inbox(host, username, password, callback, port=None, ssl=True, logger.warning("IMAP connection timeout. Reconnecting...") -def save_output(results, output_directory="output"): +def save_output(results, output_directory="output", + output_json_aggregate="aggregate.json", + output_json_forensic="forensic.json", + output_csv_aggregate="aggregate.csv", + output_csv_forensic="forensic.csv"): """ Save report data in the given directory @@ -1273,22 +1277,30 @@ def save_output(results, output_directory="output"): else: os.makedirs(output_directory) - with open("{0}".format(os.path.join(output_directory, "aggregate.json")), + with open("{0}" + .format(os.path.join(output_directory, + output_json_aggregate)), "w", newline="\n", encoding="utf-8") as agg_json: agg_json.write(json.dumps(aggregate_reports, ensure_ascii=False, indent=2)) - with open("{0}".format(os.path.join(output_directory, "aggregate.csv")), + with open("{0}" + .format(os.path.join(output_directory, + output_csv_aggregate)), "w", newline="\n", encoding="utf-8") as agg_csv: csv = parsed_aggregate_reports_to_csv(aggregate_reports) agg_csv.write(csv) - with open("{0}".format(os.path.join(output_directory, "forensic.json")), + with open("{0}" + .format(os.path.join(output_directory, + output_json_forensic)), "w", newline="\n", encoding="utf-8") as for_json: for_json.write(json.dumps(forensic_reports, ensure_ascii=False, indent=2)) - with open("{0}".format(os.path.join(output_directory, "forensic.csv")), + with open("{0}" + .format(os.path.join(output_directory, + output_csv_forensic)), "w", newline="\n", encoding="utf-8") as for_csv: csv = parsed_forensic_reports_to_csv(forensic_reports) for_csv.write(csv) diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py index 14f168b..d46e327 100644 --- a/parsedmarc/cli.py +++ b/parsedmarc/cli.py @@ -178,6 +178,18 @@ def _main(): help=strip_attachment_help, action="store_true") arg_parser.add_argument("-o", "--output", help="write output files to the given directory") + arg_parser.add_argument("--output-json-aggregate", + help="output aggregate JSON file", + default="aggregate.json") + arg_parser.add_argument("--output-json-forensic", + help="output forensic JSON file", + default="forensic.json") + arg_parser.add_argument("--output-csv-aggregate", + help="output aggregate CSV file", + default="aggregate.csv") + arg_parser.add_argument("--output-csv-forensic", + help="output forensic CSV file", + default="forensic.csv") arg_parser.add_argument("-n", "--nameservers", nargs="+", help="nameservers to query") arg_parser.add_argument("-t", "--dns_timeout", @@ -203,11 +215,16 @@ def _main(): forensic_reports = [] args = arg_parser.parse_args() + opts = Namespace(file_path=args.file_path, config_file=args.config_file, offline=args.offline, strip_attachment_payloads=args.strip_attachment_payloads, output=args.output, + output_json_aggregate=args.output_json_aggregate, + output_json_forensic=args.output_json_forensic, + output_csv_aggregate=args.output_csv_aggregate, + output_csv_forensic=args.output_csv_forensic, nameservers=args.nameservers, silent=args.silent, dns_timeout=args.dns_timeout, @@ -665,7 +682,11 @@ def _main(): ("forensic_reports", forensic_reports)]) if opts.output: - save_output(results, output_directory=opts.output) + save_output(results, output_directory=opts.output, + output_json_aggregate=opts.output_json_aggregate, + output_json_forensic=opts.output_json_forensic, + output_csv_aggregate=opts.output_csv_aggregate, + output_csv_forensic=opts.output_csv_forensic) process_reports(results)