From 83e229aeb1e563088de76c5d573bfab8a0a0ae37 Mon Sep 17 00:00:00 2001 From: atanas argirov Date: Mon, 28 Dec 2020 15:57:32 +0000 Subject: [PATCH 1/3] * added output_{json,csv}_{aggregate,forensic}_file command line args * refactored save_output() to support output_* --- parsedmarc/__init__.py | 14 +++++++++----- parsedmarc/cli.py | 14 +++++++++++++- 2 files changed, 22 insertions(+), 6 deletions(-) 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, From 36c592cc5a6cf52995763c62c6087402538946f9 Mon Sep 17 00:00:00 2001 From: atanas argirov Date: Thu, 11 Feb 2021 18:22:29 +0000 Subject: [PATCH 2/3] * added defaults for arg parser --- parsedmarc/__init__.py | 1 - parsedmarc/cli.py | 17 +++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py index 49de20b..e27aee8 100644 --- a/parsedmarc/__init__.py +++ b/parsedmarc/__init__.py @@ -1234,7 +1234,6 @@ def watch_inbox(host, username, password, callback, port=None, ssl=True, except (timeout, IMAPClientError): logger.warning("IMAP connection timeout. Reconnecting...") - def save_output(results, output_directory="output", \ output_json_aggregate_file="aggregate.json", \ output_json_forensic_file="forensic.json", \ diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py index 262a82b..b0d4d72 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: - opts.elasticsearch_hosts: + if opts.elasticsearch_hosts: shards = opts.elasticsearch_number_of_shards replicas = opts.elasticsearch_number_of_replicas elastic.save_aggregate_report_to_elasticsearch( @@ -161,13 +161,13 @@ def _main(): 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") + help="output aggregate JSON file", default="aggregate.json") arg_parser.add_argument("--output-json-forensic-file", - help="output forensic JSON file") + help="output forensic JSON file", default="forensic.json") arg_parser.add_argument("--output-csv-aggregate-file", - help="output aggregate CSV file") + help="output aggregate CSV file", default="aggregate.csv") arg_parser.add_argument("--output-csv-forensic-file", - help="output forensic CSV file") + 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", @@ -191,6 +191,7 @@ def _main(): forensic_reports = [] args = arg_parser.parse_args() + opts = Namespace(file_path=args.file_path, config_file=args.config_file, offline=args.offline, @@ -631,7 +632,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_file=opts.output_json_aggregate_file, \ + output_json_forensic_file=opts.output_json_forensic_file, \ + output_csv_aggregate_file=opts.output_csv_aggregate_file, \ + output_csv_forensic_file=opts.output_csv_forensic_file) process_reports(results) From e51f2b0127bad34903fee4c88d44b5f4660b0913 Mon Sep 17 00:00:00 2001 From: atanas argirov Date: Fri, 12 Feb 2021 10:50:25 +0000 Subject: [PATCH 3/3] * general cleanup to meet linter rules --- parsedmarc/__init__.py | 27 ++++++++++++++++++--------- parsedmarc/cli.py | 38 +++++++++++++++++++++----------------- 2 files changed, 39 insertions(+), 26 deletions(-) diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py index e27aee8..7f0c7d0 100644 --- a/parsedmarc/__init__.py +++ b/parsedmarc/__init__.py @@ -1234,11 +1234,12 @@ def watch_inbox(host, username, password, callback, port=None, ssl=True, except (timeout, IMAPClientError): logger.warning("IMAP connection timeout. Reconnecting...") -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"): + +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 @@ -1256,22 +1257,30 @@ def save_output(results, output_directory="output", \ else: os.makedirs(output_directory) - with open("{0}".format(os.path.join(output_directory, output_json_aggregate_file)), + 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, output_csv_aggregate_file)), + 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, output_json_forensic_file)), + 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, output_csv_forensic_file)), + 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 b0d4d72..777dfc5 100644 --- a/parsedmarc/cli.py +++ b/parsedmarc/cli.py @@ -160,14 +160,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-file", - help="output aggregate JSON file", default="aggregate.json") - arg_parser.add_argument("--output-json-forensic-file", - help="output forensic JSON file", default="forensic.json") - arg_parser.add_argument("--output-csv-aggregate-file", - help="output aggregate CSV file", default="aggregate.csv") - arg_parser.add_argument("--output-csv-forensic-file", - help="output forensic CSV file", default="forensic.csv") + 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", @@ -197,10 +201,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, + 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, @@ -632,11 +636,11 @@ def _main(): ("forensic_reports", forensic_reports)]) if opts.output: - save_output(results, output_directory=opts.output, \ - output_json_aggregate_file=opts.output_json_aggregate_file, \ - output_json_forensic_file=opts.output_json_forensic_file, \ - output_csv_aggregate_file=opts.output_csv_aggregate_file, \ - output_csv_forensic_file=opts.output_csv_forensic_file) + 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)