#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <stdlib.h>
#include <unistd.h>
#include <strings.h>
#include <sys/time.h>
#include <sys/socket.h>

extern const char *__progname;

static const char *cache_file = ".wi-cache";
static const char *iana_whois = "whois.iana.org";
static const char *whois_port = "43";
#define EXPIRY 2629746 /* approx one month */

static int fflag = 0;
static int qflag = 0;
static char *dom = 0;
static char *enter = 0;
static int checkflag = 0;

typedef struct ce CE;

struct ce {
  CE *link;
  char *tld;
  unsigned long int stamp;
  unsigned int status;
#define CES_NONE 1
#define CES_GOOD 2
  char *server;
  } ;

static char *tld;
static char *cache_path;
static CE *cache;
static int cache_fd;
static FILE *cf;
static struct timeval nowtv;

static void usage(void) __attribute__((__noreturn__));
static void usage(void)
{
 fprintf(stderr,"Usage: %s [options] [-dom domain] [domain]\n",__progname);
 fprintf(stderr,"A domain must be specified, one way or another\n");
 fprintf(stderr,"Options:\n");
 fprintf(stderr,"-f (`force'): always check %s\n",iana_whois);
 fprintf(stderr,"-q (`quick'): never check %s\n",iana_whois);
 fprintf(stderr,"-enter <server>: enter <server> as server for the tld\n");
 exit(1);
}

static void handleargs(int ac, char **av)
{
 int skip;
 int errs;

 skip = 0;
 errs = 0;
 for (ac--,av++;ac;ac--,av++)
  { if (skip > 0)
     { skip --;
       continue;
     }
    if (**av != '-')
     { if (! dom)
	{ dom = *av;
	}
       else
	{ fprintf(stderr,"%s: extra argument `%s'\n",__progname,*av);
	  errs ++;
	}
       continue;
     }
    if (0)
     {
needarg:;
       fprintf(stderr,"%s: %s needs a following argument\n",__progname,*av);
       errs ++;
       continue;
     }
#define WANTARG() do { if (++skip >= ac) goto needarg; } while (0)
    if (!strcmp(*av,"-f"))
     { fflag = 1;
       continue;
     }
    if (!strcmp(*av,"-q"))
     { qflag = 1;
       continue;
     }
    if (!strcmp(*av,"-enter"))
     { WANTARG();
       enter = av[skip];
       continue;
     }
    if (!strcmp(*av,"-check"))
     { checkflag = 1;
       continue;
     }
    if (!strcmp(*av,"-dom"))
     { WANTARG();
       if (dom)
	{ fprintf(stderr,"%s: domain already specified\n",__progname);
	  errs ++;
	}
       dom = av[skip];
       continue;
     }
#undef WANTARG
    fprintf(stderr,"%s: unrecognized option `%s'\n",__progname,*av);
    errs ++;
  }
 if (errs || !dom) usage();
}

static void find_tld(void)
{
 char *d;

 d = rindex(dom,'.');
 tld = d ? d+1 : dom;
 if (strlen(tld) > 63)
  { fprintf(stderr,"%s: TLD %s too long\n",__progname,tld);
    exit(1);
  }
}

static void set_cache_path(void)
{
 char *h;

 h = getenv("HOME");
 if (h == 0)
  { fprintf(stderr,"%s: no $HOME\n",__progname);
    exit(1);
  }
 asprintf(&cache_path,"%s/%s",h,cache_file);
}

static void open_cache(void)
{
 while (1)
  { cache_fd = open(cache_path,O_RDWR,0);
    if (cache_fd >= 0) break;
    if (errno != ENOENT)
     { fprintf(stderr,"%s: %s: %s\n",__progname,cache_path,strerror(errno));
       exit(1);
     }
    cache_fd = open(cache_path,O_RDWR|O_CREAT|O_TRUNC,0666);
    if (cache_fd >= 0) break;
    if (errno != EEXIST)
     { fprintf(stderr,"%s: %s: %s\n",__progname,cache_path,strerror(errno));
       exit(1);
     }
  }
 if (flock(cache_fd,LOCK_EX) < 0)
  { fprintf(stderr,"%s: flock %s: %s\n",__progname,cache_path,strerror(errno));
    exit(1);
  }
 cf = fdopen(cache_fd,"r+");
}

static void close_cache(void)
{
 flock(cache_fd,LOCK_UN);
 close(cache_fd);
}

static void load_cache(void)
{
 int c;
 CE *e;
 int l;

 rewind(cf);
 while <"entries"> (1)
  { c = getc(cf);
    if (c == EOF) break <"entries">;
    do <"partial">
     { e = malloc(sizeof(CE));
       e->status = (c & 0x80) ? CES_GOOD : CES_NONE;
       l = c & 0x3f;
       e->tld = malloc(l+1);
       if (fread(e->tld,1,l,cf) != l) break <"partial">;
       e->tld[l] = '\0';
       l = getc(cf);
       if (l == EOF) break <"partial">;
       e->stamp = 0;
       for (;l>0;l--)
	{ c = getc(cf);
	  if (c == EOF) break <"partial">;
	  e->stamp = (e->stamp << 8) + c;
	}
       switch (e->status)
	{ case CES_NONE:
	     e->server = 0;
	     break;
	  case CES_GOOD:
	     c = getc(cf);
	     if (c == EOF) break <"partial">;
	     l = getc(cf);
	     if (l == EOF) break <"partial">;
	     l |= c << 8;
	     e->server = malloc(l+1);
	     if (fread(e->server,1,l,cf) != l) break <"partial">;
	     e->server[l] = '\0';
	     break;
	}
       e->link = cache;
       cache = e;
       continue <"entries">;
     } while (0);
    fprintf(stderr,"%s: %s: dropping partial entry\n",__progname,cache_path);
    break <"entries">;
  }
}

static void save_cache(void)
{
 CE *e;
 int l;

 static void put_stamp(unsigned long int s, FILE *f)
  { int i;
    unsigned long int t;
    for (i=0,t=s;t;t>>=8,i++) ;
    putc(i,f);
    for (i--;i>=0;i--) putc((s>>(i<<3))&0xff,f);
  }

 rewind(cf);
 for (e=cache;e;e=e->link)
  { l = strlen(e->tld);
    if (l > 63) abort();
    switch (e->status)
     { case CES_GOOD:
	  putc(l|0x80,cf);
	  fwrite(e->tld,1,l,cf);
	  put_stamp(e->stamp,cf);
	  l = strlen(e->server);
	  putc(l>>8,cf);
	  putc(l&0xff,cf);
	  fwrite(e->server,1,l,cf);
	  break;
       case CES_NONE:
	  putc(l,cf);
	  fwrite(e->tld,1,l,cf);
	  put_stamp(e->stamp,cf);
	  break;
       default:
	  abort();
	  break;
     }
  }
 fflush(cf);
 ftruncate(fileno(cf),ftell(cf));
}

static CE *lookup(const char *d)
{
 CE *e;

 for (e=cache;e;e=e->link) if (!strcasecmp(e->tld,d)) break;
 return(e);
}

static int crlf_to_nl(void *fvp, char *buf, int len)
{
 int c1;
 int c2;
 int n;
 FILE *f;

 f = fvp;
 n = 0;
 for (;len>0;len--,n++)
  { c1 = getc(f);
    if (c1 == EOF) break;
    if (c1 != '\r')
     { *buf++ = c1;
       continue;
     }
    c2 = getc(f);
    if (c2 == EOF)
     { *buf++ = '\r';
       len --;
       n ++;
       break;
     }
    if (c2 != '\n')
     { *buf++ = '\r';
       ungetc(c2,f);
       continue;
     }
    *buf++ = '\n';
    continue;
  }
 return((n>0)?n:ferror(f)?-1:0);
}

static FILE *whois_stream(const char *server, const char *domain)
{
 struct addrinfo hints;
 struct addrinfo *ai0;
 struct addrinfo *ai;
 int e;
 int se;
 int s;

 hints.ai_flags = 0;
 hints.ai_family = PF_UNSPEC;
 hints.ai_socktype = SOCK_STREAM;
 hints.ai_protocol = 0;
 hints.ai_addrlen = 0;
 hints.ai_canonname = 0;
 hints.ai_addr = 0;
 hints.ai_next = 0;
 e = getaddrinfo(server,whois_port,&hints,&ai0);
 if (e)
  { fprintf(stderr,"%s: %s/%s: %s\n",__progname,server,whois_port,gai_strerror(e));
    exit(1);
  }
 if (! ai0)
  { fprintf(stderr,"%s: %s/%s: successful lookup but no addresses?\n",__progname,server,whois_port);
    exit(1);
  }
 se = 0;
 for (ai=ai0;ai;ai=ai->ai_next)
  { s = socket(ai->ai_family,ai->ai_socktype,ai->ai_protocol);
    if (s < 0)
     { if (se == 0) se = errno;
       continue;
     }
    se = -1;
    if (connect(s,ai->ai_addr,ai->ai_addrlen) < 0)
     { char hnbuf[NI_MAXHOST];
       char pnbuf[NI_MAXSERV];
       e = errno;
       if (getnameinfo(ai->ai_addr,ai->ai_addrlen,&hnbuf[0],NI_MAXHOST,&pnbuf[0],NI_MAXSERV,NI_NUMERICHOST|NI_NUMERICSERV))
	{ fprintf(stderr,"%s: %s/%s: connect failed [%s], can't get numeric hostname info [%s]\n",__progname,server,whois_port,strerror(e),strerror(errno));
	}
       else
	{ if (ai0->ai_next)
	   { fprintf(stderr,"%s: connect %s [%s/%s]: %s\n",__progname,server,&hnbuf[0],&pnbuf[0],strerror(e));
	   }
	  else
	   { fprintf(stderr,"%s: connect %s/%s: %s\n",__progname,&hnbuf[0],&pnbuf[0],strerror(e));
	   }
	}
       close(s);
       continue;
     }
    freeaddrinfo(ai0);
    write(s,domain,strlen(domain));
    write(s,"\r\n",2);
    /* Some whois servers don't return any response if we SHUT_WR the
       connection.  The protocol spec is unclear, but appears to indicate
       that the client should not send a FIN. */
    /*shutdown(s,SHUT_WR);*/
    return(fdopen(s,"r"));
  }
 if (se > 0)
  { fprintf(stderr,"%s: socket: %s\n",__progname,strerror(se));
  }
 else
  { fprintf(stderr,"%s: %s/%s: failure but no error?\n",__progname,server,whois_port);
  }
 exit(1);
}

static CE *iana_lookup(const char *d)
{
 FILE *f;
 FILE *g;
 const char *s;
 int c;
 char *b;
 int l;
 int a;
 CE *e;

 static void acc(int c)
  { if (l >= a) b = realloc(b,a=l+16);
    b[l++] = c;
  }

 printf("[Checking %s for %s...]\n",iana_whois,d);
 f = whois_stream(iana_whois,d);
 g = fropen(f,crlf_to_nl);
 b = 0;
 l = 0;
 a = 0;
 s = 0;
 while <"readwhois"> (1)
  { c = getc(g);
    switch (c)
     { case EOF:
	  break <"readwhois">;
       case '\n':
/*	  s = "Whois Server (port 43): "; */
	  s = "whois:        ";
	  break;
       default:
	  if (s)
	   { if (! *s)
	      { acc(c);
	      }
	     else if (c == *s)
	      { s ++;
		if (! *s) l = 0;
	      }
	     else
	      { s = 0;
	      }
	   }
	  break;
     }
  }
 fclose(g);
 fclose(f);
 e = malloc(sizeof(CE));
 e->stamp = nowtv.tv_sec;
 e->tld = strdup(d);
 if (b)
  { acc('\0');
    e->status = CES_GOOD;
    e->server = b;
    printf("[Got %s]\n",b);
  }
 else
  { e->status = CES_NONE;
    e->server = 0;
    printf("[No server found]\n");
  }
 return(e);
}

static CE *fake_entry(const char *d, const char *s)
{
 CE *e;

 e = malloc(sizeof(CE));
 e->stamp = nowtv.tv_sec;
 e->tld = strdup(d);
 e->status = CES_GOOD;
 e->server = strdup(s);
 return(e);
}

int main(int ac, char **av)
{
 CE *e;
 CE *n;

 handleargs(ac,av);
 find_tld();
 set_cache_path();
 open_cache();
 load_cache();
 gettimeofday(&nowtv,0);
 e = lookup(tld);
 if (enter)
  { n = fake_entry(tld,enter);
  }
 else if ((e && !fflag && (nowtv.tv_sec < e->stamp+EXPIRY)) || qflag)
  { n = 0;
  }
 else
  { n = iana_lookup(tld);
  }
 if (n)
  { if (e)
     { n->link = e->link;
       *e = *n;
     }
    else
     { n->link = cache;
       cache = n;
       e = n;
     }
    save_cache();
  }
 close_cache();
 if (! e)
  { printf("No data for %s\n",tld);
  }
 else
  { switch (e->status)
     { case CES_NONE:
	  printf("No server for %s\n",tld);
	  break;
       case CES_GOOD:
	   { FILE *f;
	     FILE *g;
	     int c;
	     if (checkflag)
	      { printf("for %s: %s\n",tld,e->server);
	      }
	     else
	      { printf("[%s]\n",e->server);
		f = whois_stream(e->server,dom);
		g = fropen(f,crlf_to_nl);
		while (1)
		 { c = getc(g);
		   if (c == EOF) break;
		   putchar(c);
		 }
		fclose(g);
		fclose(f);
	      }
	   }
	  break;
       default:
	  abort();
	  break;
     }
  }
 exit(0);
}