sniffer by pysec

/*
    changes by oaktree:

    * make argv[1] specify the number of
    packets to sniff

    * make argv[2] optional argument for
    file to write sniff results to
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <netinet/udp.h>
#include <netinet/tcp.h>
#include <netinet/ether.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <arpa/inet.h>
#include <pcap.h>

char* checkFlag(struct tcphdr* tcp)
{
    if      (tcp->th_flags == TH_SYN)
        return "SYN";
    else if (tcp->th_flags == TH_RST)
        return "RST";
    else if (tcp->th_flags == TH_FIN)
        return "FIN";

    return "Unknown";
}

int main (int argc, char** argv) {

    if (argc < 2 || argc > 3) {
        printf("Usage: ./sniffer <# of packets to sniff> [file to write to (optional)]\n");
        exit (2);
    }
    int n = atoi(argv[1]);

    /*
        initialize stuff
    */
    char err[PCAP_ERRBUF_SIZE];
    char *device;
    pcap_t *handle;

    const u_char *packet;

    if ((device = pcap_lookupdev(err)) == NULL) {
        perror("device:");
    }

    printf("\nSniffing on interface %s\n", device);

    if ((handle = pcap_open_live (device, BUFSIZ, 0, 1000, err)) == NULL)
    {
        fprintf (stderr, "%s\n", err);
        exit (1);
    }

    struct ip           *ip;
    struct udphdr       *udp;
    struct tcphdr       *tcp;
    struct ether_header *eh;
    struct icmphdr      *icmp;
    struct pcap_pkthdr header;

    int etype;

    char buf[BUFSIZ];
    FILE* fp = NULL;
    if (argc == 3) {
        fp = fopen(argv[2], "w");

        if (fp == NULL) {
            printf ("The file %s could not be opened/created.\n", argv[2]);
            exit(3);
        }
    }

    for (int i = 0; i < n; i++) {

        packet = pcap_next (handle, &header);
        usleep (0);

        if (!packet)
            continue;

        /* get some useful info */
        /* Assuming Ethernet Link 802.3 */
        eh = (struct ether_header*)packet;
        etype = ntohs (eh->ether_type);

        /* Ignore non IP packets... */
        if (etype != ETHERTYPE_IP)
            continue;

        ip = (struct ip*) (packet + ETH_HLEN);

        // clear buffer
        memset(buf, '\0', BUFSIZ);

        if (ip->ip_p == IPPROTO_UDP) {
            udp = (struct udphdr*) (packet + ETH_HLEN + sizeof(struct ip));
            sprintf (buf, "\n[ UDP ] Src Port: %u  Dst Port: %u\n", ntohs(udp->source), ntohs(udp->dest));
        }

        if (ip->ip_p == IPPROTO_TCP) {
            tcp = (struct tcphdr*) (packet + ETH_HLEN + sizeof(struct ip));
            sprintf (buf, "\n[ TCP ] Src Port: %u  Dst Port: %u Flags Set: %s\n", ntohs(tcp->source), ntohs(tcp->dest), checkFlag(tcp));
        }

        if (ip->ip_p == IPPROTO_ICMP) {
            icmp = (struct icmphdr*) (packet + ETH_HLEN + sizeof(struct ip));
            sprintf (buf, "\n[ ICMP ] Type: %u  Code: %u\n", ntohs(icmp->type), ntohs(icmp->code));
        }

        printf("%s", buf);

        if (fp != NULL) {
            fwrite(buf, sizeof(char), strlen(buf), fp);
        }

    }

    fclose (fp);
    pcap_close (handle);
    return 0;
}