# --------------------------------------------------------------------------------------------
# Copyright (c) jehama. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------

#!/usr/bin/env python3
# pylint: disable=C0301

"""Creates a universal timestamp field for hearthbleed results,
so Kibana can create a time based index of it."""
# Imports
import datetime
import json
import os
# import sys
from pathlib import Path

# A simple message to notify the user about the start of the script.
print("---Hearthbleed timestamp editor---")

# Loop through a list of files in the directory
for filename in os.listdir("./"):
    # Check if the file is a JSON file and it contains heartbleed in the name.
    if filename.endswith(".json") and "heartbleed" in filename:
        # Before editing print the name of the file that will be edited.
        
print("Editing {0}...".format(filename))

        # Make preparations with regard to the the filenames, counters and the desired date format.
        input_file = Path(filename)
        output_file = Path(filename[:-5] + "-edited.json")
        total_count = 0
        missing_count = 0
        new_date_format = "%Y-%m-%d %H:%M:%S:%f"

        # Select the file to be used or the output.
        with output_file.open("w") as outfile:
            # Loop through every JSON element in the input file.
            for element in input_file.open("r"):
                # Update a counter, reset the date field and load in the next JSON element.
                total_count += 1
                date = None
                json_element = json.loads(element)

                # If the loaded element does not contain the right json_element,
                # or if the json_element contains a different date format then a KeyError is thrown.
                try:
                    # Create a new json_element with the same value as the one already known.
                    json_element["custom.heartbleed.timestamp"] = json_element["source2.443"]["https"]["heartbleed"]["timestamp"]
                    old_date_format = "%Y-%m-%d %H:%M:%S UTC"
                    # To convert the format of the timestamp the timestamp is converted to a datetime object.
                    # Then the datetime object is converted to a string following the desired format.
                    date = datetime.datetime.strptime(json_element["custom.heartbleed.timestamp"], old_date_format)
                    new_date = datetime.datetime.strftime(date, new_date_format)
                    # The value of the new json_element is then updated with the timestamp string.
                    json_element["custom.heartbleed.timestamp"] = new_date
                except KeyError:
                    try:
                        # This json_element can contain two different timestamp formats.
                        # So a second try is used to check for the second format if the first one is incorrect.
                        json_element["custom.heartbleed.timestamp"] = json_element["source1.timestamp"]
                        old_date_format = "%Y-%m-%dT%H:%M:%S.%f"
                        try:
                            date = datetime.datetime.strptime(json_element["custom.heartbleed.timestamp"], old_date_format)
                        except ValueError:
                            date = datetime.datetime.strptime(json_element["custom.heartbleed.timestamp"], "%Y-%m-%dT%H:%M:%S")
                        new_date = datetime.datetime.strftime(date, new_date_format)
                        json_element["custom.heartbleed.timestamp"] = new_date
                    except KeyError:
                        # If these statements are reached then that means that either the timestamp formats don't match,
                        # or that there is no corresponding timestamp field in the element.
                        missing_count += 1
                        continue
                        # Instead of simply noting the failed creation of the new json_element one can also choose stop the script.
                        # This is done by uncommenting the next line while commenting two instructions above.
                        # sys.exit("Field to edit is missing or it does not conform to the given timestamp formats." + "\n Exiting now...")

                # Write updated element to the new file.
                outfile.write(json.dumps(json_element) + "\n")

                # A simple way of keeping visualising the progress, this can be uncommented if desired.
                # print("\r" + str(total_count) + " done..", end="")

        print("Total edited: {0}".format(total_count))
        print("Total missed: {0}".format(missing_count))
        print("Edited results stored in: {0}".format(output_file))