/*
 * PROJECT:     ReactOS trace route utility
 * LICENSE:     GPL-2.0-or-later (https://spdx.org/licenses/GPL-2.0-or-later)
 * PURPOSE:     Trace network paths through networks
 * COPYRIGHT:   Copyright 2018 Ged Murphy <gedmurphy@reactos.org>
 *              Copyright 2025 Curtis Wilson <LiquidFox1776@gmail.com>
 */

#include <stdarg.h>
#include <stdlib.h>

#include <windef.h>
#include <winbase.h>
#include <winuser.h>

#include <ws2tcpip.h>
#include <iphlpapi.h>
#include <icmpapi.h>
#include <winsock2.h>
#include <conutils.h>

#include <strsafe.h>

#include "resource.h"

#define SIZEOF_ICMP_ERROR       8
#define SIZEOF_IO_STATUS_BLOCK  8
#define PACKET_SIZE             32
#define MAX_IPADDRESS           32
#define NUM_OF_PINGS            3
#define MIN_HOP_COUNT           1
#define MAX_HOP_COUNT           255
#define MIN_MILLISECONDS        1
#define MAX_MILLISECONDS        ULONG_MAX

struct TraceInfo
{
    bool ResolveAddresses;
    ULONG MaxHops;
    ULONG Timeout;
    WCHAR HostName[NI_MAXHOST];
    WCHAR TargetIP[MAX_IPADDRESS];
    int Family;

    HANDLE hIcmpFile;
    PADDRINFOW Target;

} Info = { 0 };


#if 0
static
INT
OutputText(
    _In_ UINT uID,
    ...)
{
    INT Len;
    va_list args;

    va_start(args, uID);
    Len = ConResMsgPrintfV(StdOut, 0, uID, &args);
    va_end(args);

    return Len;
}
#else
#define OutputText(uID, ...) ConResMsgPrintf(StdOut, 0, (uID), ##__VA_ARGS__)
#endif

static
VOID
Usage()
{
    OutputText(IDS_USAGE);
}

static bool
GetULONG(
    _In_ PCWSTR String,
    _Out_ PULONG Value)
{
    PWSTR StopString;

    // Check input arguments
    if (*String == UNICODE_NULL)
        return false;

    // Clear errno so we can use its value after
    // the call to wcstoul to check for errors.
    errno = 0;

    // Try to convert String to ULONG
    *Value = wcstoul(String, &StopString, 10);
    if ((errno != ERANGE) && (errno != 0 || *StopString != UNICODE_NULL))
        return false;
    // The conversion was successful
    return true;
}

static bool
ResolveTarget()
{
    ADDRINFOW Hints;
    ZeroMemory(&Hints, sizeof(Hints));
    Hints.ai_family = Info.Family;
    Hints.ai_flags = AI_CANONNAME;

    int Status;
    Status = GetAddrInfoW(Info.HostName,
                          NULL,
                          &Hints,
                          &Info.Target);
    if (Status != 0)
        return false;

    Status = GetNameInfoW(Info.Target->ai_addr,
                          (socklen_t)Info.Target->ai_addrlen,
                          Info.TargetIP,
                          MAX_IPADDRESS,
                          NULL,
                          0,
                          NI_NUMERICHOST);
    return (Status == 0);
}

static bool
PrintHopInfo(_In_ PVOID Buffer)
{
    SOCKADDR_IN6 SockAddrIn6 = { 0 };
    SOCKADDR_IN SockAddrIn = { 0 };
    PSOCKADDR SockAddr;
    socklen_t Size;

    if (Info.Family == AF_INET6)
    {
        PIPV6_ADDRESS_EX Ipv6Addr = (PIPV6_ADDRESS_EX)Buffer;
        SockAddrIn6.sin6_family = AF_INET6;
        CopyMemory(SockAddrIn6.sin6_addr.u.Word, Ipv6Addr->sin6_addr, sizeof(SockAddrIn6.sin6_addr));
        //SockAddrIn6.sin6_addr = Ipv6Addr->sin6_addr;
        SockAddr = (PSOCKADDR)&SockAddrIn6;
        Size = sizeof(SockAddrIn6);

    }
    else
    {
        IPAddr *Address = (IPAddr *)Buffer;
        SockAddrIn.sin_family = AF_INET;
        SockAddrIn.sin_addr.S_un.S_addr = *Address;
        SockAddr = (PSOCKADDR)&SockAddrIn;
        Size = sizeof(SockAddrIn);
    }

    INT Status;
    bool Resolved = false;
    WCHAR HostName[NI_MAXHOST];
    if (Info.ResolveAddresses)
    {
        Status = GetNameInfoW(SockAddr,
                              Size,
                              HostName,
                              NI_MAXHOST,
                              NULL,
                              0,
                              NI_NAMEREQD);
        if (Status == 0)
        {
            Resolved = true;
        }
    }

    WCHAR IpAddress[MAX_IPADDRESS];
    Status = GetNameInfoW(SockAddr,
                          Size,
                          IpAddress,
                          MAX_IPADDRESS,
                          NULL,
                          0,
                          NI_NUMERICHOST);
    if (Status == 0)
    {
        if (Resolved)
        {
            OutputText(IDS_HOP_RES_INFO, HostName, IpAddress);
        }
        else
        {
            OutputText(IDS_HOP_IP_INFO, IpAddress);
        }
    }

    return (Status == 0);
}

static ULONG
GetResponseStats(
    _In_ PVOID ReplyBuffer,
    _Out_ ULONG& RoundTripTime,
    _Out_ PVOID& AddressInfo
)
{
    ULONG Status;

    if (Info.Family == AF_INET6)
    {
        PICMPV6_ECHO_REPLY EchoReplyV6;
        EchoReplyV6 = (PICMPV6_ECHO_REPLY)ReplyBuffer;
        Status = EchoReplyV6->Status;
        RoundTripTime = EchoReplyV6->RoundTripTime;
        AddressInfo = &EchoReplyV6->Address;
    }
    else
    {
#ifdef _WIN64
        PICMP_ECHO_REPLY32 EchoReplyV4;
        EchoReplyV4 = (PICMP_ECHO_REPLY32)ReplyBuffer;
#else
        PICMP_ECHO_REPLY EchoReplyV4;
        EchoReplyV4 = (PICMP_ECHO_REPLY)ReplyBuffer;
#endif
        Status = EchoReplyV4->Status;
        RoundTripTime = EchoReplyV4->RoundTripTime;
        AddressInfo = &EchoReplyV4->Address;
    }

    return Status;
}

static bool
DecodeResponse(
    _In_ PVOID ReplyBuffer,
    _In_ PVOID LastGoodResponse,
    _In_ bool OutputHopAddress,
    _Out_ bool& GoodResponse,
    _Out_ bool& FoundTarget
)
{
    ULONG RoundTripTime;
    PVOID AddressInfo;
    ULONG Status = GetResponseStats(ReplyBuffer, RoundTripTime, AddressInfo);

    switch (Status)
    {
    case IP_SUCCESS:
    case IP_TTL_EXPIRED_TRANSIT:
        if (RoundTripTime)
        {
            OutputText(IDS_HOP_TIME, RoundTripTime);
        }
        else
        {
            OutputText(IDS_HOP_ZERO);
        }
        GoodResponse = true;
        break;

    case IP_DEST_HOST_UNREACHABLE:
    case IP_DEST_NET_UNREACHABLE:
        FoundTarget = true;
        PrintHopInfo(AddressInfo);
        OutputText(IDS_HOP_RESPONSE);
        if (Status == IP_DEST_HOST_UNREACHABLE)
        {
            OutputText(IDS_DEST_HOST_UNREACHABLE);
        }
        else if (Status == IP_DEST_NET_UNREACHABLE)
        {
            OutputText(IDS_DEST_NET_UNREACHABLE);
        }
        return true;

    case IP_REQ_TIMED_OUT:
        OutputText(IDS_TIMEOUT);
        break;

    case IP_GENERAL_FAILURE:
        OutputText(IDS_GEN_FAILURE);
        return false;

    default:
        OutputText(IDS_TRANSMIT_FAILED, Status);
        return false;
    }

    if (OutputHopAddress)
    {
        if (Status == IP_REQ_TIMED_OUT && LastGoodResponse)
        {
            Status = GetResponseStats(LastGoodResponse, RoundTripTime, AddressInfo);
        }
        if (Status == IP_SUCCESS)
        {
            FoundTarget = true;
        }
        if (Status == IP_TTL_EXPIRED_TRANSIT || Status == IP_SUCCESS)
        {
            PrintHopInfo(AddressInfo);
            OutputText(IDS_LINEBREAK);
        }
        else if (Status == IP_REQ_TIMED_OUT)
        {
            OutputText(IDS_REQ_TIMED_OUT);
        }
    }

    return true;
}

static bool
RunTraceRoute()
{
    bool Success = false;
    PVOID ReplyBuffer = NULL, LastGoodResponse = NULL;
    DWORD ReplySize;

    HANDLE heap = GetProcessHeap();
    bool Quit = false;
    ULONG HopCount = 1;
    bool FoundTarget = false;

    Success = ResolveTarget();
    if (!Success)
    {
        OutputText(IDS_UNABLE_RESOLVE, Info.HostName);
        goto Cleanup;
    }

    ReplySize = PACKET_SIZE + SIZEOF_ICMP_ERROR + SIZEOF_IO_STATUS_BLOCK;
    if (Info.Family == AF_INET6)
    {
        ReplySize += sizeof(ICMPV6_ECHO_REPLY);
    }
    else
    {
#ifdef _WIN64
        ReplySize += sizeof(ICMP_ECHO_REPLY32);
#else
        ReplySize += sizeof(ICMP_ECHO_REPLY);
#endif
    }

    ReplyBuffer = HeapAlloc(heap, HEAP_ZERO_MEMORY, ReplySize);
    if (ReplyBuffer == NULL)
    {
        Success = false;
        goto Cleanup;
    }

    if (Info.Family == AF_INET6)
    {
        Info.hIcmpFile = Icmp6CreateFile();
    }
    else
    {
        Info.hIcmpFile = IcmpCreateFile();
    }
    if (Info.hIcmpFile == INVALID_HANDLE_VALUE)
    {
        Success = false;
        goto Cleanup;
    }

    OutputText(IDS_TRACE_INFO, Info.HostName, Info.TargetIP, Info.MaxHops);

    IP_OPTION_INFORMATION IpOptionInfo;
    ZeroMemory(&IpOptionInfo, sizeof(IpOptionInfo));

    while ((HopCount <= Info.MaxHops) && (FoundTarget == false) && (Quit == false))
    {
        OutputText(IDS_HOP_COUNT, HopCount);

        if (LastGoodResponse)
        {
            HeapFree(heap, 0, LastGoodResponse);
            LastGoodResponse = NULL;
        }

        for (int Ping = 1; Ping <= NUM_OF_PINGS; Ping++)
        {
            BYTE SendBuffer[PACKET_SIZE];
            bool GoodResponse = false;

            IpOptionInfo.Ttl = static_cast<UCHAR>(HopCount);

            if (Info.Family == AF_INET6)
            {
                struct sockaddr_in6 Source;

                ZeroMemory(&Source, sizeof(Source));
                Source.sin6_family = AF_INET6;

                (void)Icmp6SendEcho2(Info.hIcmpFile,
                                     NULL,
                                     NULL,
                                     NULL,
                                     &Source,
                                     (struct sockaddr_in6 *)Info.Target->ai_addr,
                                     SendBuffer,
                                     (USHORT)PACKET_SIZE,
                                     &IpOptionInfo,
                                     ReplyBuffer,
                                     ReplySize,
                                     Info.Timeout);
            }
            else
            {
                (void)IcmpSendEcho2(Info.hIcmpFile,
                                     NULL,
                                     NULL,
                                     NULL,
                                     ((PSOCKADDR_IN)Info.Target->ai_addr)->sin_addr.s_addr,
                                     SendBuffer,
                                     (USHORT)PACKET_SIZE,
                                     &IpOptionInfo,
                                     ReplyBuffer,
                                     ReplySize,
                                     Info.Timeout);
            }

            if (DecodeResponse(ReplyBuffer,
                               LastGoodResponse,
                               (Ping == NUM_OF_PINGS),
                               GoodResponse,
                               FoundTarget) == false)
            {
                Quit = true;
                break;
            }

            if (FoundTarget)
            {
                Success = true;
                break;
            }

            if (GoodResponse)
            {
                if (LastGoodResponse)
                {
                    HeapFree(heap, 0, LastGoodResponse);
                }
                LastGoodResponse = HeapAlloc(heap, HEAP_ZERO_MEMORY, ReplySize);
                if (LastGoodResponse == NULL)
                {
                    Success = false;
                    goto Cleanup;
                }
                CopyMemory(LastGoodResponse, ReplyBuffer, ReplySize);
            }
        }

        HopCount++;
        Sleep(100);
    }

    OutputText(IDS_TRACE_COMPLETE);

Cleanup:
    if (ReplyBuffer)
    {
        HeapFree(heap, 0, ReplyBuffer);
    }
    if (LastGoodResponse)
    {
        HeapFree(heap, 0, LastGoodResponse);
    }
    if (Info.Target)
    {
        FreeAddrInfoW(Info.Target);
    }
    if (Info.hIcmpFile)
    {
        IcmpCloseHandle(Info.hIcmpFile);
    }

    return Success;
}

static bool
GetUlongOptionInRange(
    _In_ int argc,
    _In_ wchar_t *argv[],
    _Inout_ int *i,
    _Out_ ULONG *Value,
    _In_  ULONG MinimumValue,
    _In_  ULONG MaximumValue)
{
    ULONG ParsedValue = 0;

    // Check if we have enough values
    if ((*i + 1) > (argc - 1))
    {
        OutputText(IDS_MISSING_OPTION_VALUE, argv[*i]);
        return false;
    }

    (*i)++;

    // Try to parse and convert the value as ULONG.
    // Check if ParsedValue is within the specified range.
    if (!GetULONG(argv[*i], &ParsedValue) ||
        ((ParsedValue < MinimumValue) || (ParsedValue > MaximumValue)))
    {
        (*i)--;
        OutputText(IDS_BAD_OPTION_VALUE, argv[*i]);
        return false;
    }

    *Value = ParsedValue;
    return true;
}

static bool
ParseCmdline(int argc, wchar_t *argv[])
{
    if (argc < 2)
    {
        Usage();
        return false;
    }

    for (int i = 1; i < argc; i++)
    {
        if (argv[i][0] == '-' || argv[i][0] == '/')
        {
            switch (argv[i][1])
            {
            case 'd':
                Info.ResolveAddresses = false;
                break;

            case 'h':
                if (!GetUlongOptionInRange(argc,
                                           argv,
                                           &i,
                                           &Info.MaxHops,
                                           MIN_HOP_COUNT,
                                           MAX_HOP_COUNT))
                {
                    return false;
                }
                break;

            case 'j':
                printf("-j is not yet implemented.\n");
                return false;

            case 'w':
                if (!GetUlongOptionInRange(argc,
                                           argv,
                                           &i,
                                           &Info.Timeout,
                                           MIN_MILLISECONDS,
                                           MAX_MILLISECONDS))
                {
                    return false;
                }
                break;

            case '4':
                Info.Family = AF_INET;
                break;

            case '6':
                Info.Family = AF_INET6;
                break;

            default:
            {
                OutputText(IDS_INVALID_OPTION, argv[i]);
                Usage();
                return false;
            }
            }
        }
        else
        {
            // The host must be the last argument
            if (i != (argc - 1))
            {
                Usage();
                return false;
            }
            
            StringCchCopyW(Info.HostName, NI_MAXHOST, argv[i]);
            break;
        }
    }

    // Check for missing host
    if (Info.HostName[0] == UNICODE_NULL)
    {
        OutputText(IDS_MISSING_TARGET);
        Usage();
        return false;
    }
    return true;
}

EXTERN_C
int wmain(int argc, wchar_t *argv[])
{
    /* Initialize the Console Standard Streams */
    ConInitStdStreams();

    Info.ResolveAddresses = true;
    Info.MaxHops = 30;
    Info.Timeout = 4000;
    Info.Family = AF_UNSPEC;

    if (!ParseCmdline(argc, argv))
    {
        return 1;
    }

    WSADATA WsaData;
    if (WSAStartup(MAKEWORD(2, 2), &WsaData))
    {
        return 1;
    }

    bool Success;
    Success = RunTraceRoute();

    WSACleanup();

    return Success ? 0 : 1;
}
