diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py index 4924d07..49de20b 100644 --- a/parsedmarc/__init__.py +++ b/parsedmarc/__init__.py @@ -1235,7 +1235,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_file="aggregate.json", \ + output_json_forensic_file="forensic.json", \ + output_csv_aggregate_file="aggregate.csv", \ + output_csv_forensic_file="forensic.csv"): """ Save report data in the given directory @@ -1253,22 +1257,22 @@ 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_file)), "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_file)), "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_file)), "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_file)), "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 9dad3fb..262a82b 100644 --- a/parsedmarc/cli.py +++ b/parsedmarc/cli.py @@ -82,7 +82,7 @@ def _main(): if opts.save_aggregate: for report in reports_["aggregate_reports"]: try: - if opts.elasticsearch_hosts: + opts.elasticsearch_hosts: shards = opts.elasticsearch_number_of_shards replicas = opts.elasticsearch_number_of_replicas elastic.save_aggregate_report_to_elasticsearch( @@ -160,6 +160,14 @@ 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-file", + help="output aggregate JSON file") + arg_parser.add_argument("--output-json-forensic-file", + help="output forensic JSON file") + arg_parser.add_argument("--output-csv-aggregate-file", + help="output aggregate CSV file") + arg_parser.add_argument("--output-csv-forensic-file", + help="output forensic CSV file") arg_parser.add_argument("-n", "--nameservers", nargs="+", help="nameservers to query") arg_parser.add_argument("-t", "--dns_timeout", @@ -188,6 +196,10 @@ def _main(): offline=args.offline, strip_attachment_payloads=args.strip_attachment_payloads, output=args.output, + output_json_aggregate_file=args.output_json_aggregate_file, + output_json_forensic_file=args.output_json_forensic_file, + output_csv_aggregate_file=args.output_csv_aggregate_file, + output_csv_forensic_file=args.output_csv_forensic_file, nameservers=args.nameservers, silent=args.silent, dns_timeout=args.dns_timeout,