#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <errno.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <arpa/inet.h>
#include <unistd.h>

#define BUFFER_SIZE  100  // Size of buffer for datagrams

// Datagram to send as "ping"
struct PingDatagram
{
	struct iphdr   ip;    // IP header (without options)
	struct icmphdr icmp;  // ICMP header
	struct timeval tv;    // Payload: time of sending
};

/* Name:         getCmdLineArgs
 * Description:  Gets the expected arguments from command line
 * Parameters:   argc - number of arguments specified
 *               argv - array with arguments as strings
 *               server - hostname of server to connect to
 * Return code:  true if successful
 *               false if an error occurred
 */
bool getCmdLineArgs(int argc, char *argv[], const char **server)
{
	bool result = false;

	if (argc == 2)
	{
		*server = argv[1];
		result = true;
	}
	else
	{
		fprintf(stderr, "usage: %s <server>\n", argv[0]);
	}

	return result;
}

/* Name:         resolveHostName
 * Description:  Resolves a hostname into an IP address
 * Parameters:   name - hostname to resolve
 *               addr - buffer to store IP address in
 * Return code:  true if successful
 *               false if an error occurred
 */
bool resolveHostName(const char *name, struct in_addr *addr)
{
	bool result = false;

	printf("Resolving hostname \"%s\"...", name);
	struct hostent *host = gethostbyname(name);
	if (host)
	{
		printf(" OK.\n");
		*addr = *(struct in_addr*)host->h_addr;
		result = true;
	}
	else
	{
		printf(" FAILED.\n");
		fprintf(stderr, "Couldn't resolve hostname: %s\n", hstrerror(h_errno));
	}

	return result;
}

/* Name:         createRawSocket
 * Description:  Creates a new, raw socket
 * Parameters:   sock - buffer to receive socket
 *               proto - protocol to create socket for
 * Return code:  true if successful
 *               false if an error occurred
 */
bool createRawSocket(int *sock, int proto)
{
	bool failed = false;

	*sock = socket(AF_INET, SOCK_RAW, proto);
	if (*sock == -1)
	{
		fprintf(stderr, "Couldn't create socket: %s\n", strerror(errno));
		failed = true;
	}

	return !failed;
}

/* Name:         setIpHeaderInclude
 * Description:  Set socket option IP_HDRINCL
 * Parameters:   sock - socket to set option for
 *               value - value to set option to
 * Return code:  true if successful
 *               false if an error occurred
 */
bool setIpHeaderInclude(int sock, bool value)
{
	bool failed = false;

	int opt = value;
	int result = setsockopt(sock, IPPROTO_IP, IP_HDRINCL, &opt, sizeof(opt));
	if (result != 0)
	{
		fprintf(stderr, "Setting socket option failed: %s\n", strerror(errno));
		failed = true;
	}

	return !failed;
}

/* Name:         safeCloseSocket
 * Description:  Checks if a socket is valid, closes it. The socket is
 *               set to an invalid value to avoid double free errors.
 * Parameters:   sock - socket to communicate with
 * Return code:  none
 */
void safeCloseSocket(int *sock)
{
	if (*sock != -1)
	{
		close(*sock);
		*sock = -1;
	}
}

/* Name:         computeInternetChecksum
 * Description:  Compute the Internet Checksum as described in section
 *               4.1 of RFC 1071.
 * Parameters:   addr - start address of checksum calculation
 *               count - number of bytes to include in checksum
 * Return code:  Internet Checksum of specified input data
 */
uint16_t computeInternetChecksum(const void *addr, int count)
{
	register long sum = 0;
	const uint16_t *word = addr;

	while (count > 1)
	{
		sum += *word++;
		count -= 2;
	}

	if (count > 0)
		sum += *(uint8_t*) word;

	while (sum >> 16)
		sum = (sum & 0xffff) + (sum >> 16);

	return ~sum;
}

/* Name:         sendEchoRequest
 * Description:  Send ICMP echo request with current time of day as
 *               payload using a raw socket with IP_HDRINCL set.
 * Parameters:   sock - socket to send request on
 *               addr - address to send echo request to
 * Return code:  true if successful
 *               false if an error occurred
 */
bool sendEchoRequest(int sock, struct in_addr addr)
{
	bool result = false;
	struct PingDatagram dgram;

	// IP header
	dgram.ip.ihl = sizeof(dgram.ip) / 4;  // size of IP header in 32-bit words
	dgram.ip.version = 4;                 // IPv4
	dgram.ip.tos = 0;                     // type of service
	dgram.ip.tot_len = 0;                 // total length (filled in by kernel)
	dgram.ip.id = 0;                      // datagram ID (filled in by kernel)
	dgram.ip.frag_off = 0;                // fragmentation offset
	dgram.ip.ttl = 64;                    // Time-to-live for datagram
	dgram.ip.protocol = IPPROTO_ICMP;     // Payload protocol is ICMP
	dgram.ip.check = 0;                   // checksum, filled in by kernel
	dgram.ip.saddr = 0;                   // source address, filled in by kernel
	dgram.ip.daddr = addr.s_addr;         // Destination address

	// ICMP header
	dgram.icmp.type = ICMP_ECHO;          // ICMP echo request
	dgram.icmp.code = 0;                  // sub code
	dgram.icmp.checksum = 0;              // checksum is 0 during calculation
	dgram.icmp.un.echo.id = getpid();     // ID for identifying own packets
	dgram.icmp.un.echo.sequence = 0;      // sequence number

	// ICMP payload
	gettimeofday(&dgram.tv, NULL);

	// ICMP checksum includes ICMP header and ICMP payload
	int size = sizeof(dgram.icmp) + sizeof(dgram.tv);
	dgram.icmp.checksum = computeInternetChecksum(&dgram.icmp, size);

	struct sockaddr_in saddr;
	saddr.sin_family = AF_INET;
	saddr.sin_port = 0;
	saddr.sin_addr = addr;
	int bytes = sendto(sock, &dgram, sizeof(dgram), 0, (struct sockaddr*) &saddr, sizeof(saddr));
	if (bytes > 0)
	{
		printf("Sent echo request with %i bytes to %s.\n", bytes, inet_ntoa(saddr.sin_addr));
		result = true;
	}
	else
	{
		fprintf(stderr, "Couldn't write to socket: %s\n", strerror(errno));
	}

	return result;
}

/* Name:         receiveEchoReply
 * Description:  Receive a datagram on a raw ICMP socket and check if
 *               it matches the expected echo reply.
 * Parameters:   sock - socket to receive datagram on
 * Return code:  true if successful
 *               false if an error occurred
 */
bool receiveEchoReply(int sock)
{
	bool failed = false;

	char buffer[BUFFER_SIZE];
	int bytes = recv(sock, buffer, sizeof(buffer), 0);
	if (bytes >= 0)
	{
		struct PingDatagram *dgram = (struct PingDatagram*) buffer;
		struct timeval tv;
		gettimeofday(&tv, NULL);

		struct in_addr addr = { dgram->ip.saddr };
		printf("Received %i bytes from %s.\n", bytes, inet_ntoa(addr));

		if (dgram->icmp.type == ICMP_ECHOREPLY && dgram->icmp.un.echo.id == getpid())
		{
			int ms = (tv.tv_sec - dgram->tv.tv_sec) * 1000 + (tv.tv_usec - dgram->tv.tv_usec) / 1000;
			printf("It is the reply to the echo request we sent %i ms ago.\n", ms);
		}
		else
		{
			printf("It is not the packet we are looking for. :-(\n");
		}
	}
	else
	{
		fprintf(stderr, "Couldn't read from socket: %s\n", strerror(errno));
		failed = true;
	}

	return !failed;
}

int main(int argc, char *argv[])
{
	const char *hostname;
	bool good = getCmdLineArgs(argc, argv, &hostname);

	struct in_addr addr;
	if (good)
		good = resolveHostName(hostname, &addr);

	int sock = -1;
	if (good)
		good = createRawSocket(&sock, IPPROTO_ICMP);

	/* Note: the use of IP_HDRINCL is not actually required for sending
	 * ICMP datagrams, you can simply send them directly with sendto()
	 * and let the kernel fill in the IP header. However, since this
	 * example should also show how IP_HDRINCL works, it is still used.
	 */

	if (good)
		good = setIpHeaderInclude(sock, true);

	if (good)
		good = sendEchoRequest(sock, addr);

	if (good)
		good = receiveEchoReply(sock);

	safeCloseSocket(&sock);

	return good ? EXIT_SUCCESS : EXIT_FAILURE;
}