#include <winsock2.h>
#include <Ipexport.h>
#include <process.h>
#include <stdio.h>
#include <stdint.h>
#include "pingscan.h"
#include "func.h"

/* usage()
 * No comments ;-)
 */
void usage(void)
{
    printf("Pingscan version %s\n", VERSION);
    printf("%s\n", COPYRIGHT);
    printf("Usage:\n\n");
    printf("pingscan <IP-address>[/netmask] [-n] [-s packet-size] [-c packets-count] [-w time-wait]\n\n");
    printf("Parametr \"netmask\" must be in CIDR notation: 255.255.255.0 = /24\n");
    printf("\t-n - disable reverse name resolving. Default - revers name resolving enabled.\n");
    printf("\t-s - set packet size. Default is 32 bytes.\n");
    printf("\t-c - packets count for each host. Deault is 2.\n");
    printf("\t-w - timeout waiting icmp-echo response in milliseconds. Must be >= 500,\n");
    printf("\t     default is 1000.\n");
    printf("\t-f - flood mode. Only for one host. Not for subnet! Length of subnet mask must be = 32 bits.\n");
    printf("\t     Set value of key '-c' = 0 for infinity flood.\n");
    printf("\n");
    printf("Samples.\n\n");
    printf("Scanning subnet 192.168.0.0 with netmask 255.255.255.0 without reverse names resolving:\n");
    printf("pingscan -c 3 192.168.0.1/24 -n -s 1000 -w 2000\n\n");
    printf("Infinity flood to host 192.168.1.1, packets size is 64000 bytes:\n");
    printf("pingscan.exe 192.168.1.1/32 -f -s 64000 -c 0");
    printf("\n");
    ExitProcess(0);
};

/* parsing_cmd_arguments()
 * Function for parsed command line.
 */
uint32_t parsing_cmd_arguments(uint32_t arg_count, char * cmd_args[], struct config * config_params)
{
    int32_t i;
    char * mask_ptr;
    if(arg_count < 2) 
        return 0;
    for(i = 1; i < arg_count; ++i)
    {
        if(cmd_args[i][0] == '-')
        {
            if(strlen(cmd_args[i]) > 2) 
                return 0;
            switch(cmd_args[i][1])
            {
            case 's': 
                if((i + 1) < arg_count)
                {     
                    if(sscanf(cmd_args[i + 1], "%u" , & config_params -> packet_size) != 1) return 0;
                }
                else return 0;
                if(config_params -> packet_size > 65500 || config_params -> packet_size < 0) 
                    return 0;
                ++i;
                break;
            case 'n': 
                config_params -> no_rev_resolve = 1;
                break;
            case 'c': 
                if((i + 1) < arg_count)
                {
                    if(sscanf(cmd_args[i + 1], "%u" , &config_params -> packets_count) != 1) return 0;
                }
                else return 0;
                ++i;
                break;
            case 'w': 
                if((i + 1) < arg_count)
                {
                    if(sscanf(cmd_args[i + 1], "%u" , &config_params -> wait_time) != 1) 
                        return 0;
                }
                else return 0;
                if(config_params -> wait_time < 500)
                {
                    printf("Ignored parametr -w\n");
                    config_params -> wait_time = DEF_WAIT_TIME;
                }
                ++i;
                break;
            case 'f':
                config_params -> flood = 1;
                break;
            default:
                return 0;
                break;
            }
        }
        else
        {
            if(config_params -> subnet_ptr) 
                return 0; 
            if( (mask_ptr = strstr(cmd_args[i],"/")) )
            {
                ++mask_ptr;
                *(mask_ptr - 1) = 0;
                if(inet_addr(cmd_args[i]) == INADDR_NONE)
                {
                    printf("Invalid IP address.\n");
                    return 0;
                }
                config_params -> subnet_ptr = cmd_args[i];
                if(sscanf(mask_ptr, "%u", &config_params -> netmask_len) != 1)
                {
                    printf("Invalid subnet mask.\n");
                    return 0;
                }
                else if(config_params -> netmask_len > 32 || config_params -> netmask_len < 1)
                {
                    printf("Netmask must be <= 32 and > 0\n");
                    return 0;
                }
            }
            else
            {
                if(inet_addr(cmd_args[i]) == INADDR_NONE)
                {
                    printf("Invalid IP address.\n");
                    return 0;
                }
                config_params -> subnet_ptr = cmd_args[i];
            }
        }
    }/* End for(...) */
    if(config_params -> packets_count < 1 && ! config_params -> flood) 
    {
        printf("Packets count must be > 0.\n");
        return 0;
    }
    if(!config_params -> subnet_ptr) /* Subnet address is not set from command line! */ 
        return 0;
    if(config_params -> flood && config_params -> netmask_len != 32)
        return 0;
    return 1;
}

/* one_ping()
 * Function sending ICMP echoes to one host. Execute from separate threads.
 */
unsigned __stdcall one_ping( void* arg )
{
    struct echo_params * params = (struct echo_params *)arg;
    struct in_addr from_ip;
    IP_OPTION_INFORMATION ip_info;
    int32_t i;
    char ch;
    
    ZeroMemory(&ip_info, sizeof(ip_info));
    ip_info.Ttl = 128;
    
    char *send_buffer = (char*) malloc(params -> config_params -> packet_size);
    if(! send_buffer)
    {
        fprintf(stderr, "Can't allocate memory for request ICMP buffer! Exit.\n");
        ExitProcess(-1);        
    }
    /* Filling ICMP buffer  - alphabet and special characters */
    for(i = 0, ch = '0' - 1; i < params -> config_params -> packet_size; ++i)
        send_buffer[i] = (ch++ == 126)? ch = '0': ch;
    /* Allocating memory for reply buffer */
    void * repl_buf = (void*) malloc(sizeof(ICMP_ECHO_REPLY) + params -> config_params -> packet_size);
    if(!repl_buf)
    {
        fprintf(stderr, "Can't allocate memory for reply ICMP buffer! Exit.\n");
        ExitProcess(-1);
    }
    /* Send ICMP echo */
    HANDLE h_icmp_file;
    uint8_t * hostname = 0;
    PICMP_ECHO_REPLY icmpEcho = 0;
    h_icmp_file = icmp_create_file();
    for(i = 0; i < params -> config_params -> packets_count; ++i)
    {
        icmp_echo(
                h_icmp_file,            /* Handle from IcmpCreateFile() */
                params -> dst_ip,        /* Destination IP address */
                send_buffer,            /* Pointer to buffer to send */
                params -> config_params -> packet_size,    /* Size of buffer in bytes */
                &ip_info,                /* Request options */
                repl_buf,                /* Reply buffer */
                params -> config_params -> packet_size + sizeof(ICMP_ECHO_REPLY),
                params -> config_params -> wait_time - 500);    /* Time to wait in milliseconds. 500 - is default
                                                                     built-in value */
        icmpEcho = (PICMP_ECHO_REPLY)repl_buf;
        from_ip.s_addr = icmpEcho -> Address;

        if((icmpEcho -> Status) == IP_SUCCESS)
        {
            if(icmpEcho->Address == params -> dst_ip)
            {
                hostname = 0;
                if(!params -> config_params -> no_rev_resolve)
                    hostname = rev_ns_resolve(icmpEcho->Address, 0, 0);
                if(hostname)
                {
                    printf("Found %-15s %-25s time=%ld ms TTL=%d size=%u bytes\n",
                                            inet_ntoa(from_ip), hostname,
                                            icmpEcho->RoundTripTime,
                                            icmpEcho->Options.Ttl, icmpEcho->DataSize);
                }
                else
                {
                    printf("Found %-15s time=%ld ms TTL=%d size=%u bytes\n",
                                            inet_ntoa(from_ip),
                                            icmpEcho->RoundTripTime,
                                            icmpEcho->Options.Ttl,icmpEcho->DataSize);
                }
                EnterCriticalSection(params -> host_incr);
                ++ (*params -> hosts_count);
                LeaveCriticalSection(params -> host_incr);
            }
            break;    
        }
    }
    free(send_buffer);
    free(repl_buf);
    if(!icmp_close_file(h_icmp_file))
    {
        printf("Failed close_icmp_file!\n");
        _endthreadex(-1);
        return -1;
    }
    
    if(params -> release_mem_after_use)
        free(params);
    _endthreadex(0);
    return 0;
}
/* keyboard_interrupt_handler()
 * Ctrl + C handler. 
 */
BOOL WINAPI keyboard_interrupt_handler(DWORD event)
{
    end_scan_time = GetTickCount();
    if(event == CTRL_C_EVENT)
    {
        if(flood_mode)
        {
            printf(" (%u packets)\n", unfinished_echoes_count);
            printf("Program terminated at %g seconds.\n", (float)(end_scan_time - start_scan_time)/1000);        
            printf("Flood statistic: %u packets transmitted, %u received, %.2f%% packets loss.\n",
                    total_sended, total_sended - unfinished_echoes_count,
                    ((float)unfinished_echoes_count / total_sended) * 100);
        }
        else
        {
            printf("\nProgram terminated at %g seconds. ", (float)(end_scan_time - start_scan_time)/1000);        
            printf("%d host(s) found.\n", hosts_count);
        }
        return 0;
    }
    return 0;
}

/* rev_ns_resolve()
 * Function for reverse name resolving.
 */
int8_t * rev_ns_resolve(uint32_t ip_addr, char * hostname, uint32_t hostname_len)
{
    struct in_addr addr;
    struct hostent * host_info;
    addr.S_un.S_addr = ip_addr;
    host_info = gethostbyaddr((void*) &addr, sizeof(addr), AF_INET);
    if(!host_info || !host_info->h_name)
        return 0;
    ZeroMemory(hostname, hostname_len);
    if(hostname)
        strncpy(hostname, host_info -> h_name, hostname_len);
    return host_info -> h_name;
}

/* safe_inc_dec()
 * Safe incrementing or decremented variable from different threads.
 */
inline int32_t safe_inc_dec(CRITICAL_SECTION * critical_section, int32_t * counter, uint8_t inc_dec)
{
    EnterCriticalSection(critical_section);
    (inc_dec)? ++ *counter: -- *counter;
    LeaveCriticalSection(critical_section);
    return *counter;
}