/* This code is placed into public domain */
/* Code created by Radim Kolar hsn@netmag.cz 2004
   for more programs come to:
   http://home.tiscali.cz/~cz210552/
*/
/* Of course, I am a very interrested in user feedback!! */

/* Starring:
	     Classic unix forkbomb
	     Memory hog generator
	     CPU hog generator
	     Realloc tester
	     House of dead
*/

#include <signal.h>
#include <stdio.h>
#include <getopt.h>
#include <stdlib.h>
#include <unistd.h>

#define DEFAULT_ALARM 300
#define VERSION "1.4"
#define DEFAULT_INCPAGES 1024
#define DEFAULT_MAXMEMORY 1024*1024*2047

static int cpuhang;
static int runasroot;
static int touchmemory;
static int allocmemory;
static int zombie;
static int safetyalarm=1;
static int forkbomb;
static int quit;
static int maxtime=DEFAULT_ALARM;
static int incpages=DEFAULT_INCPAGES;
static size_t maxmemory=DEFAULT_MAXMEMORY;

static const struct option long_options[]=
{
 {"cpuhang",no_argument,&cpuhang,1},
 {"forkbomb",no_argument,&forkbomb,1},
 {"zombie",no_argument,&zombie,1},
 {"allocmemory",no_argument,&allocmemory,1},
 {"touchmemory",no_argument,&touchmemory,1},
 {"runasroot",no_argument,&runasroot,1},
 {"quit",no_argument,&quit,1},
 {"nosafetyalarm",no_argument,&safetyalarm,0},
 {"infinite",no_argument,&safetyalarm,0},
 {"maxtime",required_argument,NULL,'t'},
 {"memorylimit",required_argument,NULL,'l'},
 {"incpages",required_argument,NULL,'i'},
 {"version",no_argument,NULL,'V'},
 {"help",no_argument,NULL,'?'},
 {NULL,0,NULL,0}
};

static void fbusage(void)
{
   fprintf(stderr,
	"forkbomb [option]...\n"
	"  -c|--cpuhang             Eat CPU cycles.\n"
	"  -f|--forkbomb            Create a load of processes.\n"
	"  -z|--zombie              Bring us to House of Dead.\n"
	"  -m|--allocmemory         Allocate all available memory.\n"
	"  -M|--touchmemory         Touch allocated memory (implies -m).\n"
        "  -t|--maxtime <sec>       Stop after x seconds.\n"
        "  -l|--memorylimit <MB>    Do not use more than x MB per process.\n"
        "  -i|--incpages <num>      Allocate memory by groups of x pages.\n"
        "  --nosafetyalarm|--infinite Infinite run.\n"
        "  --runasroot              Allow running as root.\n"
        "  --quit                   Don't wait for safety timer expiration.\n"
	"  -?|-h|--help             This information.\n"
	"  -V|--version             Display program version.\n"
	);
};

static void kill_childrens(int signal)
{
  printf("Safety alarm!\n"); 
  kill(0,SIGINT);
  exit(0);
}

int main(int argc, char *argv[])
{
   int opt=0;
   int options_index=0;
   int rc=0;
   char *mem;
   char *newmem;
   size_t allocated=0;
   int pagesize;
   int i;
   char c;

   /* parse cmd junk */
   while((opt=getopt_long(argc,argv,"czmMfh?Vt:l:i:",long_options,&options_index))!=EOF )
   {
      switch(opt)
      {
       case 'c': cpuhang=1;break;
       case 'z': zombie=1;break;
       case 'M': touchmemory=1;
       case 'm': allocmemory=1;break; 
       case 'f': forkbomb=1;break;
       case 'l': maxmemory=atoi(optarg)*1024*1024;break;
       case 'i': incpages=atoi(optarg);break;
       case 'h':
       case '?': fbusage();return 0;
       case 'V':
                 printf(VERSION"\n");
                 return 0;
       case 't':
                 maxtime=atoi(optarg);
                 break;	     
      }
   }
  
   if(cpuhang==0 && allocmemory==0 && forkbomb==0 && zombie==0)
   {
       printf("Nothing to do! Use -h for help or read man page.\n");
       return 1;
   }

   if(!runasroot && geteuid()==0)
   {
      printf("I refuse to run as root. Use -h for help.\n");
      return 1;
   }

   if(zombie && forkbomb)
   {
       printf("Bombing zombies by forkbombs is not permited by law.\n");
       return 1;
   }

   if(safetyalarm)
   {
       signal(SIGALRM,kill_childrens);
       alarm(maxtime);
       printf("Safety alarm at %d sec. enabled.\n",maxtime);
   }

   pagesize=getpagesize();
   if(touchmemory) allocmemory=1;
   
   printf("Actions:");
   if(forkbomb) printf(" forkbomb");
   if(zombie) printf(" zombies");
   if(allocmemory) printf(" alloc %d MB (step %d kB)",maxmemory/1024/1024,incpages*pagesize/1024);
   if(touchmemory) printf(" and touch it");
   if(cpuhang) printf(" eat CPU time");
   printf(".\n");

   printf("Forkbomb "VERSION" started.\n");
      /*  fun begins */

      /* make childrens first */
      if(forkbomb || zombie)
      {
          while(rc>=0)
	  {
              rc=fork();
	      if(zombie && rc==0) _exit(0);
	  }
      }
      /* alloc memory */
      mem=NULL;
      if(allocmemory)
      {
           while( 
                  newmem=realloc(mem,allocated+pagesize*incpages)
                )
           {
	        if(touchmemory)
		{
		    /* touch newly allocated memory */
		    for(i=allocated;i<allocated+pagesize*incpages;i+=pagesize)
		    {
			newmem[i]=c++;
		    }
		}
                allocated+=pagesize*incpages;
                mem=newmem;
                if(allocated>=maxmemory)
                     break;
           }
      }

      if(touchmemory && cpuhang)
      {
         c=0;
         while(1)
         {
	     for(i=0;i<allocated;i+=pagesize)
		 mem[i]=c++;
	 }
      }

      if(cpuhang)
      {
         while(1); 
      } else
      { 
            if(!quit) 
	    {
		while(1)
		{
	           sleep(DEFAULT_ALARM);
		}
	    }
      } 

    return 0;
}