src/flowtop.c

   1 /*
   2  * netsniff-ng - the packet sniffing beast
   3  * By Daniel Borkmann <daniel@netsniff-ng.org>
   4  * Copyright 2011 Daniel Borkmann.
   5  * Subject to the GPL, version 2.
   6  *
   7  * A tiny tool to provide top-like netfilter connection tracking information.
   8  * Regarding locking the current code is just broken. It also needs clean ups
   9  * in general.
  10  *
  11  * Debian: apt-get install libnetfilter-conntrack3 libnetfilter-conntrack-dev
  12  *                         liburcu0 liburcu-dev
  13  *
  14  * Start conntrack (if not yet running):
  15  *   iptables -A INPUT -p tcp -m state --state ESTABLISHED -j ACCEPT
  16  *   iptables -A OUTPUT -p tcp -m state --state NEW,ESTABLISHED -j ACCEPT
  17  */
  18
  19 #define _LGPL_SOURCE
  20 #include <stdio.h>
  21 #include <stdint.h>
  22 #include <stdlib.h>
  23 #include <signal.h>
  24 #include <getopt.h>
  25 #include <pthread.h>
  26 #include <signal.h>
  27 #include <netdb.h>
  28 #include <ctype.h>
  29 #include <libnetfilter_conntrack/libnetfilter_conntrack.h>
  30 #include <libnetfilter_conntrack/libnetfilter_conntrack_tcp.h>
  31 #include <GeoIP.h>
  32 #include <GeoIPCity.h>
  33 #include <netinet/in.h>
  34 #include <curses.h>
  35 #include <dirent.h>
  36 #include <sys/stat.h>
  37 #include <urcu.h>
  38 #include <libgen.h>
  39
  40 #include "die.h"
  41 #include "xmalloc.h"
  42 #include "compiler.h"
  43 #include "misc.h"
  44 #include "signals.h"
  45 #include "locking.h"
  46 #include "timespec.h"
  47 #include "dissector_eth.h"
  48
  49 #define INCLUDE_UDP     (1 << 0)
  50 #define INCLUDE_TCP     (1 << 1)
  51
  52 #ifndef ATTR_TIMESTAMP_START
  53 # define ATTR_TIMESTAMP_START 63
  54 #endif
  55 #ifndef ATTR_TIMESTAMP_STOP
  56 # define ATTR_TIMESTAMP_STOP 64
  57 #endif
  58
  59 #define SCROLL_MAX 1000
  60
  61 struct flow_entry {
  62         uint32_t flow_id;
  63         int first;
  64         struct flow_entry *next;
  65         uint32_t use;
  66         uint32_t status;
  67         uint8_t  l3_proto;
  68         uint8_t  l4_proto;
  69         uint32_t ip4_src_addr;
  70         uint32_t ip4_dst_addr;
  71         uint32_t ip6_src_addr[4];
  72         uint32_t ip6_dst_addr[4];
  73         uint16_t port_src;
  74         uint16_t port_dst;
  75         uint8_t  tcp_state;
  76         uint8_t  tcp_flags;
  77         uint64_t counter_pkts;
  78         uint64_t counter_bytes;
  79         uint64_t timestamp_start;
  80         uint64_t timestamp_stop;
  81         char country_src[128];
  82         char city_src[128];
  83         char rev_dns_src[256];
  84         char country_dst[128];
  85         char city_dst[128];
  86         char rev_dns_dst[256];
  87         int procnum;
  88         int inode;
  89         char cmdline[256];
  90 };
  91
  92 struct flow_list {
  93         struct flow_entry *head;
  94         struct spinlock lock;
  95 };
  96
  97 static sig_atomic_t sigint = 0;
  98
  99 static double interval = 0.1;
 100
 101 static int what = INCLUDE_TCP;
 102
 103 static struct flow_list flow_list;
 104
 105 static GeoIP *gi_country = NULL;
 106 static GeoIP *gi_city = NULL;
 107
 108 static char *path_city_db = NULL, *path_country_db = NULL;
 109
 110 static const char *short_options = "t:vhTULK";
 111
 112 static struct option long_options[] = {
 113         {"interval", required_argument, 0, 't'},
 114         {"tcp", no_argument, 0, 'T'},
 115         {"udp", no_argument, 0, 'U'},
 116         {"city-db", required_argument, 0, 'L'},
 117         {"country-db", required_argument, 0, 'K'},
 118         {"version", no_argument, 0, 'v'},
 119         {"help", no_argument, 0, 'h'},
 120         {0, 0, 0, 0}
 121 };
 122
 123 const char *const l3proto2str[AF_MAX] = {
 124         [AF_INET]                       = "ipv4",
 125         [AF_INET6]                      = "ipv6",
 126 };
 127
 128 const char *const proto2str[IPPROTO_MAX] = {
 129         [IPPROTO_TCP]                   = "tcp",
 130         [IPPROTO_UDP]                   = "udp",
 131         [IPPROTO_UDPLITE]               = "udplite",
 132         [IPPROTO_ICMP]                  = "icmp",
 133         [IPPROTO_ICMPV6]                = "icmpv6",
 134         [IPPROTO_SCTP]                  = "sctp",
 135         [IPPROTO_GRE]                   = "gre",
 136         [IPPROTO_DCCP]                  = "dccp",
 137         [IPPROTO_IGMP]                  = "igmp",
 138         [IPPROTO_IPIP]                  = "ipip",
 139         [IPPROTO_EGP]                   = "egp",
 140         [IPPROTO_PUP]                   = "pup",
 141         [IPPROTO_IDP]                   = "idp",
 142         [IPPROTO_RSVP]                  = "rsvp",
 143         [IPPROTO_IPV6]                  = "ip6tun",
 144         [IPPROTO_ESP]                   = "esp",
 145         [IPPROTO_AH]                    = "ah",
 146         [IPPROTO_PIM]                   = "pim",
 147         [IPPROTO_COMP]                  = "comp",
 148 };
 149
 150 const char *const state2str[TCP_CONNTRACK_MAX] = {
 151         [TCP_CONNTRACK_NONE]            = "NOSTATE",
 152         [TCP_CONNTRACK_SYN_SENT]        = "SYN_SENT",
 153         [TCP_CONNTRACK_SYN_RECV]        = "SYN_RECV",
 154         [TCP_CONNTRACK_ESTABLISHED]     = "ESTABLISHED",
 155         [TCP_CONNTRACK_FIN_WAIT]        = "FIN_WAIT",
 156         [TCP_CONNTRACK_CLOSE_WAIT]      = "CLOSE_WAIT",
 157         [TCP_CONNTRACK_LAST_ACK]        = "LAST_ACK",
 158         [TCP_CONNTRACK_TIME_WAIT]       = "TIME_WAIT",
 159         [TCP_CONNTRACK_CLOSE]           = "CLOSE",
 160         [TCP_CONNTRACK_SYN_SENT2]       = "SYN_SENT2",
 161 };
 162
 163 const uint8_t states[] = {
 164         TCP_CONNTRACK_SYN_SENT,
 165         TCP_CONNTRACK_SYN_RECV,
 166         TCP_CONNTRACK_ESTABLISHED,
 167         TCP_CONNTRACK_FIN_WAIT,
 168         TCP_CONNTRACK_CLOSE_WAIT,
 169         TCP_CONNTRACK_LAST_ACK,
 170         TCP_CONNTRACK_TIME_WAIT,
 171         TCP_CONNTRACK_CLOSE,
 172         TCP_CONNTRACK_SYN_SENT2,
 173         TCP_CONNTRACK_NONE,
 174 };
 175
 176 static void signal_handler(int number)
 177 {
 178         switch (number) {
 179         case SIGINT:
 180                 sigint = 1;
 181                 break;
 182         case SIGHUP:
 183         default:
 184                 break;
 185         }
 186 }
 187
 188 static void help(void)
 189 {
 190         printf("\nflowtop %s, top-like netfilter TCP/UDP flow tracking\n",
 191                VERSION_STRING);
 192         printf("http://www.netsniff-ng.org\n\n");
 193         printf("Usage: flowtop [options]\n");
 194         printf("Options:\n");
 195         printf("  -t|--interval <time>   Refresh time in sec (default 0.1)\n");
 196         printf("  -T|--tcp               Show only TCP flows (default)\n");
 197         printf("  -U|--udp               Show only UDP flows\n");
 198         printf("  --city-db <path>       Specifiy path for geoip city database\n");
 199         printf("  --country-db <path>    Specifiy path for geoip country database\n");
 200         printf("  -v|--version           Print version\n");
 201         printf("  -h|--help              Print this help\n");
 202         printf("\n");
 203         printf("Examples:\n");
 204         printf("  flowtop -U --interval 0.5\n");
 205         printf("  flowtop\n\n");
 206         printf("Note:\n");
 207         printf("  If netfilter is not running, you can activate it with i.e.:\n");
 208         printf("   iptables -A INPUT -p tcp -m state --state ESTABLISHED -j ACCEPT\n");
 209         printf("   iptables -A OUTPUT -p tcp -m state --state NEW,ESTABLISHED -j ACCEPT\n");
 210         printf("\n");
 211         printf("Please report bugs to <bugs@netsniff-ng.org>\n");
 212         printf("Copyright (C) 2011 Daniel Borkmann <daniel@netsniff-ng.org>\n");
 213         printf("License: GNU GPL version 2\n");
 214         printf("This is free software: you are free to change and redistribute it.\n");
 215         printf("There is NO WARRANTY, to the extent permitted by law.\n\n");
 216         die();
 217 }
 218
 219 static void version(void)
 220 {
 221         printf("\nflowtop %s, top-like netfilter TCP/UDP flow tracking\n",
 222                VERSION_STRING);
 223         printf("http://www.netsniff-ng.org\n\n");
 224         printf("Please report bugs to <bugs@netsniff-ng.org>\n");
 225         printf("Copyright (C) 2011 Daniel Borkmann\n");
 226         printf("License: GNU GPL version 2\n");
 227         printf("This is free software: you are free to change and redistribute it.\n");
 228         printf("There is NO WARRANTY, to the extent permitted by law.\n\n");
 229         die();
 230 }
 231
 232 static void screen_init(WINDOW **screen)
 233 {
 234         (*screen) = initscr();
 235         noecho();
 236         cbreak();
 237         keypad(stdscr, TRUE);
 238         nodelay(*screen, TRUE);
 239         refresh();
 240         wrefresh(*screen);
 241 }
 242
 243 static inline uint16_t get_port(uint16_t src, uint16_t dst)
 244 {
 245         char *tmp1, *tmp2;
 246
 247         src = ntohs(src);
 248         dst = ntohs(dst);
 249
 250         /* XXX: Is there a better way to determine? */
 251         if (src < dst && src < 1024) {
 252                 return src;
 253         } else if (dst < src && dst < 1024) {
 254                 return dst;
 255         } else {
 256                 tmp1 = lookup_port_tcp(src);
 257                 tmp2 = lookup_port_tcp(dst);
 258                 if (tmp1 && !tmp2) {
 259                         return src;
 260                 } else if (!tmp1 && tmp2) {
 261                         return dst;
 262                 } else {
 263                         if (src < dst)
 264                                 return src;
 265                         else
 266                                 return dst;
 267                 }
 268         }
 269 }
 270
 271 static void screen_update(WINDOW *screen, struct flow_list *fl, int skip_lines)
 272 {
 273         int i, line = 3;
 274         int maxx, maxy;
 275         struct flow_entry *n;
 276
 277         curs_set(0);
 278         getmaxyx(screen, maxy, maxx);
 279
 280         start_color();
 281         init_pair(1, COLOR_RED, COLOR_BLACK);
 282         init_pair(2, COLOR_BLUE, COLOR_BLACK);
 283         init_pair(3, COLOR_YELLOW, COLOR_BLACK);
 284         init_pair(4, COLOR_GREEN, COLOR_BLACK);
 285
 286         clear();
 287
 288         rcu_read_lock();
 289
 290         mvwprintw(screen, 1, 2, "Kernel netfilter TCP/UDP flow statistics, [+%d] t=%.2lfs",
 291                   skip_lines, interval);
 292
 293         if (rcu_dereference(fl->head) == NULL)
 294                 mvwprintw(screen, line, 2, "(No active sessions! Is netfilter running?)");
 295
 296         maxy -= 4;
 297         /* Yes, that's lame :-P */
 298         for (i = 0; i < sizeof(states); i++) {
 299                 n = rcu_dereference(fl->head);
 300
 301                 while (n && maxy > 0) {
 302                         char tmp[128];
 303
 304                         if (n->tcp_state != states[i] ||
 305                             (i != TCP_CONNTRACK_NONE &&
 306                              n->tcp_state == TCP_CONNTRACK_NONE) ||
 307                             /* Filter out DNS */
 308                             get_port(n->port_src, n->port_dst) == 53) {
 309                                 n = rcu_dereference(n->next);
 310                                 continue;
 311                         }
 312
 313                         if (skip_lines > 0) {
 314                                 n = rcu_dereference(n->next);
 315                                 skip_lines--;
 316                                 continue;
 317                         }
 318
 319                         snprintf(tmp, sizeof(tmp), "%u/%s", n->procnum,
 320                                  basename(n->cmdline));
 321                         tmp[sizeof(tmp) - 1] = 0;
 322
 323                         mvwprintw(screen, line, 2, "[");
 324                         attron(COLOR_PAIR(3));
 325                         printw("%s", n->procnum > 0 ? tmp : "bridged(?)");
 326                         attroff(COLOR_PAIR(3));
 327                         printw("]:%s:%s[", l3proto2str[n->l3_proto],
 328                                proto2str[n->l4_proto]);
 329                         attron(COLOR_PAIR(3));
 330                         printw("%s", state2str[n->tcp_state]);
 331                         attroff(COLOR_PAIR(3));
 332                         printw("]:");
 333                         attron(A_BOLD);
 334                         if (n->tcp_state != TCP_CONNTRACK_NONE) {
 335                                 printw("%s -> ", lookup_port_tcp(get_port(n->port_src,
 336                                                                           n->port_dst)));
 337                         } else {
 338                                 printw("%s -> ", lookup_port_udp(get_port(n->port_src,
 339                                                                           n->port_dst)));
 340                         }
 341                         attroff(A_BOLD);
 342                         attron(COLOR_PAIR(1));
 343                         printw("%s", n->rev_dns_src);
 344                         attroff(COLOR_PAIR(1));
 345                         printw(":%u (", ntohs(n->port_src));
 346                         attron(COLOR_PAIR(4));
 347                         printw("%s", (strlen(n->country_src) > 0 ?
 348                                n->country_src : "N/A"));
 349                         attroff(COLOR_PAIR(4));
 350                         printw(", %s) => ", (strlen(n->city_src) > 0 ?
 351                                n->city_src : "N/A"));
 352                         attron(COLOR_PAIR(2));
 353                         printw("%s", n->rev_dns_dst);
 354                         attroff(COLOR_PAIR(2));
 355                         printw(":%u (", ntohs(n->port_dst));
 356                         attron(COLOR_PAIR(4));
 357                         printw("%s", strlen(n->country_dst) > 0 ?
 358                                n->country_dst : "N/A");
 359                         attroff(COLOR_PAIR(4));
 360                         printw(", %s)", strlen(n->city_dst) > 0 ?
 361                                n->city_dst : "N/A");
 362
 363                         line++;
 364                         maxy--;
 365                         n = rcu_dereference(n->next);
 366                 }
 367         }
 368
 369         rcu_read_unlock();
 370
 371         wrefresh(screen);
 372         refresh();
 373 }
 374
 375 static void screen_end(void)
 376 {
 377         endwin();
 378 }
 379
 380 static void presenter(void)
 381 {
 382         int skip_lines = 0;
 383         WINDOW *screen = NULL;
 384
 385         dissector_init_ethernet(0);
 386         screen_init(&screen);
 387         rcu_register_thread();
 388
 389         while (!sigint) {
 390                 switch (getch()) {
 391                 case 'q':
 392                         sigint = 1;
 393                         break;
 394                 case KEY_UP:
 395                 case 'u':
 396                 case 'k':
 397                         skip_lines--;
 398                         if (skip_lines < 0)
 399                                 skip_lines = 0;
 400                         break;
 401                 case KEY_DOWN:
 402                 case 'd':
 403                 case 'j':
 404                         skip_lines++;
 405                         if (skip_lines > SCROLL_MAX)
 406                                 skip_lines = SCROLL_MAX;
 407                         break;
 408                 default:
 409                         fflush(stdin);
 410                         break;
 411                 }
 412
 413                 screen_update(screen, &flow_list, skip_lines);
 414                 xnanosleep(interval);
 415         }
 416
 417         rcu_unregister_thread();
 418         screen_end();
 419         dissector_cleanup_ethernet();
 420 }
 421
 422 static inline const char *make_n_a(const char *p)
 423 {
 424         return p ? : "N/A";
 425 }
 426
 427 static void walk_process(char *process, struct flow_entry *n)
 428 {
 429         DIR *dir;
 430         struct dirent *ent;
 431         char path[1024];
 432
 433         if (snprintf(path, sizeof(path), "/proc/%s/fd", process) == -1)
 434                 panic("giant process name! %s\n", process);
 435
 436         dir = opendir(path);
 437         if (!dir)
 438                 return;
 439
 440         while ((ent = readdir(dir))) {
 441                 struct stat statbuf;
 442
 443                 if (snprintf(path, sizeof(path), "/proc/%s/fd/%s",
 444                              process, ent->d_name) < 0)
 445                         continue;
 446                 if (stat(path, &statbuf) < 0)
 447                         continue;
 448                 if (S_ISSOCK(statbuf.st_mode) && n->inode == statbuf.st_ino) {
 449                         memset(n->cmdline, 0, sizeof(n->cmdline));
 450                         snprintf(path, sizeof(path), "/proc/%s/exe", process);
 451                         readlink(path, n->cmdline, sizeof(n->cmdline) - 1);
 452                         n->procnum = atoi(process);
 453                 }
 454         }
 455
 456         closedir(dir);
 457 }
 458
 459 /* Derived from ifpromisc, Fred N. van Kempen, GPL v2.0 */
 460 /* n->inode must be set */
 461 static void walk_processes(struct flow_entry *n)
 462 {
 463         DIR *dir;
 464         struct dirent *ent;
 465
 466         if (n->inode <= 0) {
 467                 memset(n->cmdline, 0, sizeof(n->cmdline));
 468                 return;
 469         }
 470
 471         dir = opendir("/proc");
 472         if (!dir)
 473                 panic("Cannot open /proc!\n");
 474
 475         while ((ent = readdir(dir)))
 476                 if (strspn(ent->d_name, "0123456789") == strlen(ent->d_name))
 477                         walk_process(ent->d_name, n);
 478
 479         closedir(dir);
 480 }
 481
 482 static int get_inode_from_local_port(int port, const char *proto, int ip6)
 483 {
 484         int ret = -ENOENT;
 485         char path[128];
 486         char buff[1024];
 487         FILE *proc;
 488
 489         memset(path, 0, sizeof(path));
 490         snprintf(path, sizeof(path), "/proc/net/%s%s", proto, ip6 ? "6" : "");
 491         proc = fopen(path, "r");
 492         if (!proc)
 493                 return -EIO;
 494         memset(buff, 0, sizeof(buff));
 495         while (fgets(buff, sizeof(buff), proc) != NULL) {
 496                 int lport = 0, inode = 0;
 497                 buff[sizeof(buff) - 1] = 0;
 498                 if (sscanf(buff, "%*u: %*X:%X %*X:%*X %*X %*X:%*X %*X:%*X "
 499                            "%*X %*u %*u %u", &lport, &inode) == 2) {
 500                         if (lport == port) {
 501                                 ret = inode;
 502                                 break;
 503                         }
 504                 }
 505                 memset(buff, 0, sizeof(buff));
 506         }
 507         fclose(proc);
 508
 509         return ret;
 510 }
 511
 512 static void flow_entry_from_ct(struct flow_entry *n, struct nf_conntrack *ct)
 513 {
 514         n->flow_id = nfct_get_attr_u32(ct, ATTR_ID);
 515         n->use = nfct_get_attr_u32(ct, ATTR_USE);
 516         n->status = nfct_get_attr_u32(ct, ATTR_STATUS);
 517         n->l3_proto = nfct_get_attr_u8(ct, ATTR_ORIG_L3PROTO);
 518         n->l4_proto = nfct_get_attr_u8(ct, ATTR_ORIG_L4PROTO);
 519         n->ip4_src_addr = nfct_get_attr_u32(ct, ATTR_ORIG_IPV4_SRC);
 520         n->ip4_dst_addr = nfct_get_attr_u32(ct, ATTR_ORIG_IPV4_DST);
 521
 522         const uint8_t *ipv6_src = nfct_get_attr(ct, ATTR_ORIG_IPV6_SRC);
 523         if (ipv6_src)
 524                 memcpy(n->ip6_src_addr, ipv6_src, sizeof(n->ip6_src_addr));
 525         const uint8_t *ipv6_dst = nfct_get_attr(ct, ATTR_ORIG_IPV6_DST);
 526         if (ipv6_dst)
 527                 memcpy(n->ip6_dst_addr, ipv6_dst, sizeof(n->ip6_dst_addr));
 528
 529         n->port_src = nfct_get_attr_u16(ct, ATTR_ORIG_PORT_SRC);
 530         n->port_dst = nfct_get_attr_u16(ct, ATTR_ORIG_PORT_DST);
 531         n->tcp_state = nfct_get_attr_u8(ct, ATTR_TCP_STATE);
 532         n->tcp_flags = nfct_get_attr_u8(ct, ATTR_TCP_FLAGS_ORIG);
 533         n->counter_pkts = nfct_get_attr_u64(ct, ATTR_ORIG_COUNTER_PACKETS);
 534         n->counter_bytes = nfct_get_attr_u64(ct, ATTR_ORIG_COUNTER_BYTES);
 535         n->timestamp_start = nfct_get_attr_u64(ct, ATTR_TIMESTAMP_START);
 536         n->timestamp_stop = nfct_get_attr_u64(ct, ATTR_TIMESTAMP_STOP);
 537
 538         if (n->first) {
 539                 n->inode = get_inode_from_local_port(ntohs(n->port_src),
 540                                                      proto2str[n->l4_proto],
 541                                                      !!(ipv6_src));
 542                 if (n->inode > 0)
 543                         walk_processes(n);
 544         }
 545         /* if this really runs on a router, we try it once and then let it be */
 546         n->first = 0;
 547 }
 548
 549 /* TODO: IP4 + IP6 */
 550 static void flow_entry_get_extended(struct flow_entry *n)
 551 {
 552         struct sockaddr_in sa;
 553         struct hostent *hent;
 554         GeoIPRecord *gir_src, *gir_dst;
 555
 556         if (n->flow_id == 0)
 557                 return;
 558         if (ntohs(n->port_src) == 53 || ntohs(n->port_dst) == 53)
 559                 return;
 560
 561         memset(&sa, 0, sizeof(sa));
 562         sa.sin_family = PF_INET; //XXX: IPv4
 563         sa.sin_addr.s_addr = n->ip4_src_addr;
 564         getnameinfo((struct sockaddr *) &sa, sizeof(sa), n->rev_dns_src,
 565                     sizeof(n->rev_dns_src), NULL, 0, NI_NUMERICHOST);
 566
 567         hent = gethostbyaddr(&sa.sin_addr, sizeof(sa.sin_addr), PF_INET);
 568         if (hent) {
 569                 memset(n->rev_dns_src, 0, sizeof(n->rev_dns_src));
 570                 memcpy(n->rev_dns_src, hent->h_name,
 571                        min(sizeof(n->rev_dns_src), strlen(hent->h_name)));
 572         }
 573
 574         gir_src = GeoIP_record_by_ipnum(gi_city, ntohl(n->ip4_src_addr));
 575         if (gir_src) {
 576                 const char *country =
 577                         make_n_a(GeoIP_country_name_by_ipnum(gi_country,
 578                                                              ntohl(n->ip4_src_addr)));
 579                 const char *city = make_n_a(gir_src->city);
 580                 memcpy(n->country_src, country,
 581                        min(sizeof(n->country_src), strlen(country)));
 582                 memcpy(n->city_src, city,
 583                        min(sizeof(n->city_src), strlen(city)));
 584         }
 585
 586         memset(&sa, 0, sizeof(sa));
 587         sa.sin_family = PF_INET; //XXX: IPv4
 588         sa.sin_addr.s_addr = n->ip4_dst_addr;
 589         getnameinfo((struct sockaddr *) &sa, sizeof(sa), n->rev_dns_dst,
 590                     sizeof(n->rev_dns_dst), NULL, 0, NI_NUMERICHOST);
 591
 592         hent = gethostbyaddr(&sa.sin_addr, sizeof(sa.sin_addr), PF_INET);
 593         if (hent) {
 594                 memset(n->rev_dns_dst, 0, sizeof(n->rev_dns_dst));
 595                 memcpy(n->rev_dns_dst, hent->h_name,
 596                        min(sizeof(n->rev_dns_dst), strlen(hent->h_name)));
 597         }
 598
 599         gir_dst = GeoIP_record_by_ipnum(gi_city, ntohl(n->ip4_dst_addr));
 600         if (gir_dst) {
 601                 const char *country =
 602                         make_n_a(GeoIP_country_name_by_ipnum(gi_country,
 603                                                              ntohl(n->ip4_dst_addr)));
 604                 const char *city = make_n_a(gir_dst->city);
 605                 memcpy(n->country_dst, country,
 606                        min(sizeof(n->country_dst), strlen(country)));
 607                 memcpy(n->city_dst, city,
 608                        min(sizeof(n->city_dst), strlen(city)));
 609         }
 610 }
 611
 612 static void flow_list_init(struct flow_list *fl)
 613 {
 614         fl->head = NULL;
 615         spinlock_init(&fl->lock);
 616 }
 617
 618 static struct flow_entry *__flow_list_find_by_id(struct flow_list *fl, uint32_t id)
 619 {
 620         struct flow_entry *n = rcu_dereference(fl->head);
 621         while (n != NULL) {
 622                 if (n->flow_id == id)
 623                         return n;
 624                 n = rcu_dereference(n->next);
 625         }
 626         return NULL;
 627 }
 628
 629 static struct flow_entry *__flow_list_find_prev_by_id(struct flow_list *fl, uint32_t id)
 630 {
 631         struct flow_entry *n = rcu_dereference(fl->head);
 632         if (n->flow_id == id)
 633                 return NULL;
 634         while (rcu_dereference(n->next) != NULL) {
 635                 if (rcu_dereference(n->next)->flow_id == id)
 636                         return n;
 637                 n = rcu_dereference(n->next);
 638         }
 639         return NULL;
 640 }
 641
 642 static void flow_list_new_entry(struct flow_list *fl, struct nf_conntrack *ct)
 643 {
 644         struct flow_entry *n = xzmalloc(sizeof(*n));
 645         n->first = 1;
 646         rcu_assign_pointer(n->next, fl->head);
 647         rcu_assign_pointer(fl->head, n);
 648         flow_entry_from_ct(n, ct);
 649         flow_entry_get_extended(n);
 650 }
 651
 652 static void flow_list_update_entry(struct flow_list *fl, struct nf_conntrack *ct)
 653 {
 654         int do_ext = 0;
 655         uint32_t id = nfct_get_attr_u32(ct, ATTR_ID);
 656         struct flow_entry *n;
 657         n = __flow_list_find_by_id(fl, id);
 658         if (n == NULL) {
 659                 n = xzmalloc(sizeof(*n));
 660                 n->first = 1;
 661                 rcu_assign_pointer(n->next, fl->head);
 662                 rcu_assign_pointer(fl->head, n);
 663                 do_ext = 1;
 664         }
 665         flow_entry_from_ct(n, ct);
 666         if (do_ext)
 667                 flow_entry_get_extended(n);
 668 }
 669
 670 static void flow_list_destroy_entry(struct flow_list *fl, struct nf_conntrack *ct)
 671 {
 672         uint32_t id = nfct_get_attr_u32(ct, ATTR_ID);
 673         struct flow_entry *n1, *n2;
 674
 675         n1 = __flow_list_find_by_id(fl, id);
 676         if (n1) {
 677                 n2 = __flow_list_find_prev_by_id(fl, id);
 678                 if (n2) {
 679                         rcu_assign_pointer(n2->next, n1->next);
 680                         rcu_assign_pointer(n1->next, NULL);
 681                         xfree(n1);
 682                 } else {
 683                         xfree(fl->head);
 684                         rcu_assign_pointer(fl->head, NULL);
 685                 }
 686         }
 687 }
 688
 689 static void flow_list_destroy(struct flow_list *fl)
 690 {
 691         struct flow_entry *n;
 692
 693         while (fl->head != NULL) {
 694                 n = rcu_dereference(fl->head->next);
 695                 rcu_assign_pointer(fl->head->next, NULL);
 696                 xfree(fl->head);
 697                 rcu_assign_pointer(fl->head, n);
 698         }
 699
 700         synchronize_rcu();
 701         spinlock_destroy(&fl->lock);
 702 }
 703
 704 static int collector_cb(enum nf_conntrack_msg_type type,
 705                         struct nf_conntrack *ct,
 706                         void *data)
 707 {
 708         if (sigint)
 709                 return NFCT_CB_STOP;
 710
 711         synchronize_rcu();
 712
 713         spinlock_lock(&flow_list.lock);
 714         switch (type) {
 715         case NFCT_T_NEW:
 716                 flow_list_new_entry(&flow_list, ct);
 717                 break;
 718         case NFCT_T_UPDATE:
 719                 flow_list_update_entry(&flow_list, ct);
 720                 break;
 721         case NFCT_T_DESTROY:
 722                 flow_list_destroy_entry(&flow_list, ct);
 723                 break;
 724         default:
 725                 break;
 726         }
 727         spinlock_unlock(&flow_list.lock);
 728
 729         return NFCT_CB_CONTINUE;
 730 }
 731
 732 static int dummy_cb(enum nf_conntrack_msg_type type, struct nf_conntrack *ct,
 733                     void *data)
 734 {
 735         return NFCT_CB_STOP;
 736 }
 737
 738 static void *collector(void *null)
 739 {
 740         int ret;
 741         u_int32_t family = AF_INET;
 742         struct nfct_handle *handle;
 743         struct nfct_filter *filter;
 744
 745         handle = nfct_open(CONNTRACK, NFCT_ALL_CT_GROUPS);
 746         if (!handle)
 747                 panic("Cannot create a nfct handle!\n");
 748
 749         /* Hack: inits ct */
 750         nfct_callback_register(handle, NFCT_T_ALL, dummy_cb, NULL);
 751         nfct_query(handle, NFCT_Q_DUMP, &family);
 752         nfct_close(handle);
 753
 754         handle = nfct_open(CONNTRACK, NFCT_ALL_CT_GROUPS);
 755         if (!handle)
 756                 panic("Cannot create a nfct handle!\n");
 757
 758         nfct_query(handle, NFCT_Q_FLUSH, &family);
 759
 760         filter = nfct_filter_create();
 761         if (!filter)
 762                 panic("Cannot create a nfct filter!\n");
 763         if (what & INCLUDE_UDP)
 764                 nfct_filter_add_attr_u32(filter, NFCT_FILTER_L4PROTO, IPPROTO_UDP);
 765         if (what & INCLUDE_TCP)
 766                 nfct_filter_add_attr_u32(filter, NFCT_FILTER_L4PROTO, IPPROTO_TCP);
 767
 768         struct nfct_filter_ipv4 filter_ipv4 = {
 769                 .addr = ntohl(INADDR_LOOPBACK),
 770                 .mask = 0xffffffff,
 771         };
 772
 773         nfct_filter_set_logic(filter, NFCT_FILTER_SRC_IPV4,
 774                               NFCT_FILTER_LOGIC_NEGATIVE);
 775         nfct_filter_add_attr(filter, NFCT_FILTER_SRC_IPV4, &filter_ipv4);
 776
 777         struct nfct_filter_ipv6 filter_ipv6 = {
 778                 .addr = { 0x0, 0x0, 0x0, 0x1 },
 779                 .mask = { 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff },
 780         };
 781
 782         nfct_filter_set_logic(filter, NFCT_FILTER_SRC_IPV6,
 783                               NFCT_FILTER_LOGIC_NEGATIVE);
 784         nfct_filter_add_attr(filter, NFCT_FILTER_SRC_IPV6, &filter_ipv6);
 785
 786         ret = nfct_filter_attach(nfct_fd(handle), filter);
 787         if (ret < 0)
 788                 panic("Cannot attach filter to handle!\n");
 789
 790         nfct_filter_destroy(filter);
 791
 792         if (path_country_db)
 793                 gi_country = GeoIP_open(path_country_db, GEOIP_MMAP_CACHE);
 794         else
 795                 gi_country = GeoIP_open_type(GEOIP_COUNTRY_EDITION,
 796                                              GEOIP_MMAP_CACHE);
 797
 798         if (path_city_db)
 799                 gi_city = GeoIP_open(path_city_db, GEOIP_MMAP_CACHE);
 800         else
 801                 gi_city = GeoIP_open_type(GEOIP_CITY_EDITION_REV1,
 802                                           GEOIP_MMAP_CACHE);
 803         if (!gi_country || !gi_city)
 804                 panic("Cannot open GeoIP database!\n");
 805
 806         GeoIP_set_charset(gi_country, GEOIP_CHARSET_UTF8);
 807         GeoIP_set_charset(gi_city, GEOIP_CHARSET_UTF8);
 808
 809         flow_list_init(&flow_list);
 810
 811         rcu_register_thread();
 812
 813         nfct_callback_register(handle, NFCT_T_ALL, collector_cb, NULL);
 814
 815         while (!sigint)
 816                 nfct_catch(handle);
 817
 818         rcu_unregister_thread();
 819
 820         flow_list_destroy(&flow_list);
 821
 822         GeoIP_delete(gi_city);
 823         GeoIP_delete(gi_country);
 824
 825         nfct_close(handle);
 826
 827         if (path_city_db)
 828                 xfree(path_city_db);
 829         if (path_country_db)
 830                 xfree(path_country_db);
 831
 832         pthread_exit(0);
 833 }
 834
 835 int main(int argc, char **argv)
 836 {
 837         pthread_t tid;
 838         int ret, c, opt_index, what_cmd = 0;
 839
 840         check_for_root_maybe_die();
 841
 842         while ((c = getopt_long(argc, argv, short_options, long_options,
 843                &opt_index)) != EOF) {
 844                 switch (c) {
 845                 case 't':
 846                         if (!optarg)
 847                                 help();
 848                         interval = atof(optarg);
 849                         if (interval < 0.01)
 850                                 panic("Choose larger interval!\n");
 851                         break;
 852                 case 'T':
 853                         what_cmd |= INCLUDE_TCP;
 854                         break;
 855                 case 'U':
 856                         what_cmd |= INCLUDE_UDP;
 857                         break;
 858                 case 'L':
 859                         path_city_db = xstrdup(optarg);
 860                         break;
 861                 case 'K':
 862                         path_country_db = xstrdup(optarg);
 863                         break;
 864                 case 'h':
 865                         help();
 866                         break;
 867                 case 'v':
 868                         version();
 869                         break;
 870                 case '?':
 871                         switch (optopt) {
 872                         case 'L':
 873                         case 'K':
 874                                 panic("Option -%c requires an argument!\n",
 875                                       optopt);
 876                         default:
 877                                 if (isprint(optopt))
 878                                         whine("Unknown option character "
 879                                               "`0x%X\'!\n", optopt);
 880                                 die();
 881                         }
 882                 default:
 883                         break;
 884                 }
 885         }
 886
 887         if (what_cmd > 0)
 888                 what = what_cmd;
 889
 890         rcu_init();
 891
 892         register_signal(SIGINT, signal_handler);
 893         register_signal(SIGHUP, signal_handler);
 894
 895         ret = pthread_create(&tid, NULL, collector, NULL);
 896         if (ret < 0)
 897                 panic("Cannot create phthread!\n");
 898
 899         presenter();
 900         return 0;
 901 }
 902