/*
 * Copyright (C) 1994,1995 Lars Fenneberg
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 */

/* Parts extracted from Net-3, Authors Fred N. van Kempen, Alan Cox,
 * and from route.c, Authors Fred N. van Kempen, Johannes Stille, Linus 
 * Torvalds, Alan Cox. 
 */

#include <sys/types.h>
#include <sys/ioctl.h>
#include <ctype.h>
#include <pwd.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <netinet/in.h>

#include "pathnames.h"
#include <linux/ip_acct_user.h>

#define DENY 1
#define ALLOW 0

char pname[255];
int raw = 0;
int disp_weight = 0;
int no_header = 0;
int skfd = -1;

void usage(void)
{
  fprintf(stderr,"Usage: %s [-wvVrahCn] [-c <user>] [-d <user>] [-g <user>]\n",pname);
  fprintf(stderr,"       %s add <addr> [mask <addr>] [weight <int>]\n",pname);
  fprintf(stderr,"       %s add default [weight <int>]\n",pname);
  fprintf(stderr,"       %s del <addr> [mask <addr>]\n",pname);
  exit(1);
}

void version(void)
{
  fprintf(stderr,"%s: version %s\n",pname, VERSION);
  exit(1);
}

void drv_version(void)
{
  fprintf(stderr,"%s: compiled for kernel driver: %s", pname, IPAU_VERSION);
  exit(1);
}

void error(char *str)
{
  fprintf(stderr,"%s: %s.\n",pname, str);
  exit(1);
}

int is_addr(char *str)
{
  while(isdigit(*str)) str++;
  if (*(str++) != '.') return 0; 
  while(isdigit(*str)) str++;
  if (*(str++) != '.') return 0; 
  while(isdigit(*str)) str++;
  if (*(str++) != '.') return 0; 
  while(isdigit(*str)) str++;
  if (!isdigit(*str) && (*str !=  '\0')) return 0;
  return 1;	
}

/*
 *	Display an IP address in readable format. 
 */
 
char *in_ntoa(unsigned long in)
{
	static char buff[18];
	char *p;

	p = (char *) &in;
	sprintf(buff, "%d.%d.%d.%d",
		(p[0] & 255), (p[1] & 255), (p[2] & 255), (p[3] & 255));
	return(buff);
}


/*
 *	Convert an ASCII string to binary IP. 
 */
 
unsigned long in_aton(char *str)
{
	unsigned long l;
	unsigned int val;
	int i;

	l = 0;
	for (i = 0; i < 4; i++) 
	{
		l <<= 8;
		if (*str != '\0') 
		{
			val = 0;
			while (*str != '\0' && *str != '.') 
			{
				val *= 10;
				val += *str - '0';
				str++;
			}
			l |= val;
			if (*str != '\0') 
				str++;
		}
	}
	return(htonl(l));
}


unsigned long in_xton(char *str)
{

	char *sp = str;
	char *ptr;
	unsigned long saddr;
	int i,val;

	ptr = (unsigned char *) (&(saddr) + 1);
	for (i = 0; i < sizeof(saddr); i++) {
		val = 0;
		if (*sp == '\t')
			break;
		if (*sp >= 'A')
			val = (int) (*sp - 'A') + 10;
		else
			val = (int) (*sp - '0');
		val <<= 4;
		sp++;
		if (*sp >= 'A')
			val |= (int) (*sp - 'A') + 10;
		else
			val |= (int) (*sp - '0');
		*--ptr = (unsigned char) (val & 0377);
		sp++;
	}
	return saddr;
}

unsigned long default_mask(unsigned long dst)
{
	dst = ntohl(dst);
	if ((dst & 0xff) != 0) return htonl(0xffffffff);
	if (IN_CLASSA(dst))
		return htonl(IN_CLASSA_NET);
	if (IN_CLASSB(dst))
		return htonl(IN_CLASSB_NET);
	return htonl(IN_CLASSC_NET);
}

int reset_counter(char *user)
{
	int uid;
	struct passwd *pw;
	
	if (user && ((pw = getpwnam(user)) || (pw = getpwuid(atoi(user)))))
	{
		uid = pw->pw_uid;
	}
	else
	{
	  if (user) 
	  {
	    fprintf(stderr,"%s: unknown user %s.\n",pname, user);
	    return (1);
	  } 
	  else
	  {
	    uid = IPAU_CLEAR_ALL;
	  }
	}

	if (ioctl(skfd, SIOCRSTUT, &uid) < 0) {
		fprintf(stderr, "%s: SIOCRSTUT: %s.\n",pname, strerror(errno));
		return (1);
	}
	return (0);
}

int ctrl_access(char *user, int deny)
{
	struct ip_acct_user_ctrl_entry ct;
	struct passwd *pw;
	
	if ((pw = getpwnam(user)) || (pw = getpwuid(atoi(user))))
	{
		ct.uid = pw->pw_uid;
		ct.flags = deny?IPAU_NO_ACCESS:0;
	}
	else
	{
		fprintf(stderr,"%s: unknown user %s.\n",pname, user);
	        return (1);
	}

	if (ioctl(skfd, SIOCCTRLU, &ct) < 0) {
		fprintf(stderr, "%s: SIOCCTRLU: %s.\n",pname, strerror(errno));
		return (1);
	}
	return (0);
}

int weight_add(unsigned long waddr, unsigned long wmask, int weight)
{
	struct ip_acct_weight_entry wt;
	
	wt.addr = waddr;
	wt.mask = wmask;
	wt.weight = weight;

	if (ioctl(skfd, SIOCADDWT, &wt) < 0) {
		fprintf(stderr, "%s: SIOCADDWT: %s.\n",pname, strerror(errno));
		return (1);
	}
	return (0);
} 

int weight_del(unsigned long waddr, unsigned long wmask)
{
	struct ip_acct_weight_entry wt;
	
	wt.addr = waddr;
	wt.mask = wmask;

	if (ioctl(skfd, SIOCDELWT, &wt) < 0) {
		fprintf(stderr, "%s: SIOCDELWT: %s.\n",pname, strerror(errno));
		return (1);
	}
	return (0);
} 

void read_table(void)
{
	FILE *table;
	char waddr[18], wmask[18], buff[1024];
	int weight;
	int line = 0;
	
	if ((table = fopen(_PATH_ETC_IP_ACCT_USER,"r")) == NULL)
	{
	  perror(_PATH_ETC_IP_ACCT_USER);
	}
	
	while (fgets(buff, 1023, table)) 
	{
	  line ++;
	  if (buff[0] == '#') continue;
	
	  if (sscanf(buff, "%17s %17s %i\n",waddr, wmask, &weight)==3)
	  {
	    if (!is_addr(waddr) || !is_addr(wmask))
	    {
	      fprintf(stderr, "%s: wrong address format (line %i).\n", pname, line);
	    }
	    if (weight<0) 
	    {
	      fprintf(stderr, "%s: weight < 0 (line %i).\n", pname, line);
	    }
	    weight_add(in_aton(waddr),in_aton(wmask),weight);	  
	  } 
	  else fprintf(stderr, "%s: syntax error (line %i).\n", pname, line);
	}
	
	fclose(table);
	exit(0);
}

void print_wt(void)
{
  
  char buff[1024], waddr[18], wmask[18];
  int weight,num;
  FILE *fp;
  
  if ((fp = fopen(_PATH_PROCNET_IP_ACCT_WEIGHT, "r")) == NULL) 
  {
     perror(_PATH_PROCNET_IP_ACCT_WEIGHT);
     return;
  }
   
   if (!no_header) printf("Addr            Mask            Weight\n");
   
   while (fgets(buff,1023,fp)) 
   {
   
     num = sscanf(buff, "%s %s %i\n",waddr,wmask,&weight);
     if (num != 3) continue;

     if (!raw)
     {
       if (in_xton(waddr)==htonl(0x00000000)) 
       {
       	 strcpy(waddr,"default");
       	 strcpy(wmask,"*");
       }
       else
       {
         strcpy(waddr,in_ntoa(in_xton(waddr)));
         strcpy(wmask,in_ntoa(in_xton(wmask)));   
       }
     }
     printf("%-15s %-15s %-3i\n",waddr,wmask,weight);
   }
   fclose(fp);
}

void print_ut(void)
{
  
  char buff[1024], sflags[5];
  int sent, recv, num, uid, flags;
  FILE *fp;
  struct passwd *pw;
  
  if ((fp = fopen(_PATH_PROCNET_IP_ACCT_USER, "r")) == NULL) 
  {
     perror(_PATH_PROCNET_IP_ACCT_USER);
     return;
  }
   
   if (!no_header) printf("User       Sent       Received   Flags\n");
   
   while (fgets(buff,1023,fp)) 
   {
   
     num = sscanf(buff, "%i %i %i %x\n",&uid,&sent,&recv,&flags);
     if (num != 4)  
     {
       num = sscanf(buff, "%i %i %i\n",&uid,&sent,&recv);
       if (num!=3) continue;
       flags = 0;
     }
     
     if (sent == 0 && recv == 0 && flags == 0) continue;

     *sflags = '\0';
     strcat(sflags, (flags & IPAU_NO_ACCESS)?"N":"");

     if (raw || (pw = getpwuid((uid_t)uid)) == NULL) 
     {
        printf("%-10i %-10i %-10i %-4s\n",uid,sent,recv,sflags);	
     }
     else
     {
        printf("%-10s %-10i %-10i %-4s\n",pw->pw_name,sent,recv,sflags);
     	free(pw);
     }
   }
   fclose(fp);
}

int main(int argc, char **argv)
{
  extern int optind; 
  extern char *optarg;
  int c, n, nmax, add, def;
  int weight = IPAU_DEF_WEIGHT;
  unsigned long waddr = 0;
  unsigned long wmask = 0;

  if (strrchr(argv[0],'/') == NULL) 
  {
    strcpy(pname,argv[0]);
  }
  else
  {
    strcpy(pname,strrchr(argv[0],'/')+1);
  }

  if ((skfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) 
  {
    fprintf(stderr,"%s: can't create socket (%s)\n.",pname, strerror(errno));
    return 1;
  }

  while ((c = getopt(argc,argv,"wvVrahc:Cnd:g:")) > 0) 
  {
  	switch (c)
  	{
  		case 'w':
  				disp_weight++;
  				break;
		case 'v':
				version();
				break;
		case 'V':	
				drv_version();
				break;
		case 'r':	
				raw++;
				break;
		case 'a':	
				read_table();
				break;
		case 'c':
				return (reset_counter(optarg));
		case 'C':
				return (reset_counter(NULL));
		case 'd':
				return (ctrl_access(optarg,DENY));
		case 'g':
				return (ctrl_access(optarg,ALLOW));
		case 'n':
				no_header++;
				break;
  		case 'h':
  		default:
  				usage();
  				break; /* never reached */
  	}
  }

  n = optind;
  nmax = argc-1;

  if (n>nmax) 
  {
    if (disp_weight) print_wt(); else print_ut();
  } 
  else
  {  
    if (!strcasecmp(argv[n],"add")) 
    {
      add = 1;
    }
    else if (!strcasecmp(argv[n],"del"))
    {
      add = 0;
    }
    else error("wrong command (use 'add' or 'del')");
    
    n++;
    if (n>nmax) error("address missing");

    if (add && !strcasecmp(argv[n],"default"))
    {
      waddr = htonl(0x00000000);
      wmask = htonl(0xffffffff);
      def = 1;
    }
    else
    {
      if (!is_addr(argv[n])) error("address has wrong format");
      waddr = in_aton(argv[n]);
      def = 0;
    }

    n++;	
    while (n<=nmax)
    {
       if (!def && !strcasecmp(argv[n],"mask")) 
       {
	 n++; if (n>nmax) error("attribute 'mask' requires an argument");
	 if (!is_addr(argv[n])) error("attribute 'mask': address has wrong format");
	 wmask = in_aton(argv[n]);
       }
       else if (add && !strcasecmp(argv[n],"weight"))
       {
	 n++; if (n>nmax) error("attribute 'weight' requires an argument");
	 weight = atoi(argv[n]);
       } else error((add)?((def)?"wrong attribute (use 'weight')":"wrong attribute (use 'mask' and 'weight')"):
       			  "wrong attribute (use 'mask')");
       n++;
    }

    if (wmask==0) wmask = default_mask(waddr);
    if (add) return weight_add(waddr, wmask, weight); else return weight_del(waddr,wmask);
  
  }

  return 0;
}
