"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Friday, 28th Jun, 2024

    OBJECTIVE:

        To provide a way to ping a server and get the traceroute dump.

    REFERENCES:

        1) https://www.geeksforgeeks.org/traceroute-implementation-on-python/

    DOWNLOADS:

        N/A

"""

#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# Utils:
from utils_v2.string import json
from utils_v2.string import regex
from utils_v2.date_time import date_time

# For networking:
import socket
from scapy.all import *

# For running the script from the terminal:
import argparse


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


def get_name_and_addr(destination):

    # If the user provided the IP address:
    if (
            regex.match(destination, regex.REGEX_IPV4) or
            regex.match(destination, regex.REGEX_IPV6)
    ):
        try: destination_name = socket.gethostbyaddr(destination)[0]
        except Exception as exception: destination_name = "*"
        destination_ip = destination

    # If the provided destination was the domain name:
    else:
        destination_name = destination
        try: destination_ip = socket.gethostbyname(destination)
        except Exception as exception: destination_ip = "*"

    # Done here:
    return destination_name, destination_ip


# ---------------------------------------------------------------------------------------------------------------------


def tracert(
        destination,
        max_hops = 30,
        timeout = 2.0,
        port = 33434
):

    # Initialize the variables:
    destination_name, destination_ip = get_name_and_addr(destination)
    full_trace = []
    ttl = 1

    # Keep noting hops till the limit is reached:
    while ttl <= max_hops:

        # Create a JSON for this stage:
        this_hop = {
            "destAddr": destination_ip,
            "destName": destination_name,
            "hopNo": ttl - 1,
            "isDest": False,
            "hopAddr": None,
            "hopName": None,
            "ts": None
        }

        # Create the IP and UDP headers and combine them:
        ip_packet = IP(dst = destination, ttl = ttl)
        udp_packet = UDP(dport = port)
        trace_packet = ip_packet / udp_packet

        # Send the packet and receive a reply and note down the timestamp:
        reply = sr1(trace_packet, timeout = timeout, verbose = 0)
        this_hop["ts"] = date_time.get_current_utc_date_time(as_string = True)

        # No response:
        if reply is None: this_hop["hopAddr"] = this_hop["hopName"] = "*"

        # If some response was received, we note the values and break out if this was the destination hop:
        else:
            this_hop["hopName"], this_hop["hopAddr"] = get_name_and_addr(f"{reply.src}")
            if reply.type == 3:
                this_hop["isDest"] = True
                full_trace.append(this_hop)
                break

        # Carry on to the next hop:
        full_trace.append(this_hop)
        ttl += 1

    # Done here:
    return full_trace


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description = "Traceroute Implementation in Python!")

    parser.add_argument(
        "dest",
        help = "Destination (Name or IP address)."
    )
    parser.add_argument(
        "-m",
        "--max-hops",
        type = int,
        default = 30,
        help =
        "Maximum number of hops (default: 30)."
    )
    parser.add_argument(
        "-t",
        "--timeout",
        type = float,
        default = 2.0,
        help = "Timeout for each packet in seconds (default: 2.0)."
    )
    parser.add_argument(
        "-p",
        "--port",
        type = int,
        default = 33434,
        help = "The port to probe at the destination (default: 33434)."
    )

    args = parser.parse_args()

    trace = tracert(
        destination = args.dest,
        max_hops = args.max_hops,
        timeout = args.timeout,
        port = args.port
    )
    print("TRACE:", json.to_string(trace))