#include <stdio.h>
#include <stdlib.h>

/*
 * Assumptions:  When a machine is being infected, it will not respond
 * to probes.
 * I'm going to be lazy and only do the second (scan) phase.
 * The hit list phase is just going to assume sequential/limited rate
 * and be added to the initial time.

 * This is a variation version which just uses rc5_64 as a PRNG with the
 * lower 32 bits being the result returned.  Lazy and use the same PRNG for
 * all.  Probably not as good, but hey.
 */

/* Definitions.  Make sure these work! */
typedef unsigned int unsigned32;
typedef unsigned short unsigned16;
typedef unsigned long long unsigned64;

/*
 * All random generation is going to be using RC5-32/12/32 
 * (32 bit block size, 12 rounds, 32 bit key size (making it bigger doesn't
 * really change things).
 *
 * There is one permutation for the state of the net, and a DIFFERENT
 * permutation that the worms use for infection pattern.  When the worm
 * needs to pick a "New" random target, the target index in the permutation
 * is selected as being E(MyIP + targetIp) to create the index.  This isn't
 * quite "random", but actually is somewhat better, because it means that
 * any two worms, when they hit the same already-infected target, will jump
 * to DIFFERENT portions in the permutation!
 */

/* Code for RC5 from description in Applied Cryptography.
 */

/*
  #define numRounds 12
*/

/* Reduced roundcount to speed up running time */
#define numRounds 6

#define numRounds32 12


#define ROTR16(x,c) ((unsigned16) (((x)>>((c) & 0xf))|((x)<<(16-((c) & 0xf)))))
#define ROTL16(x,c) ((unsigned16) (((x)<<((c) & 0xf))|((x)>>(16-((c) & 0xf)))))

#define ROTR32(x,c) ((unsigned32) (((x)>>((c) & 0x1f))|((x)<<(32-((c) & 0x1f)))))
#define ROTL32(x,c) ((unsigned32) (((x)<<((c) & 0x1f))|((x)>>(32-((c) & 0x1f)))))


unsigned16 *rc5_keygen(unsigned32 key){
  unsigned16 *keyArray;
  unsigned16 lArray[2];
  int i = 0;
  int j = 0;
  int k = 0;
  unsigned16 A = 0;
  unsigned16 B = 0;
  lArray[0] = key;
  lArray[1] = key >> 16;
  keyArray = (unsigned16 *) malloc(sizeof(unsigned32) * 2 * (numRounds + 1)); 
  keyArray[0] = 0xb7e5; 
  for(i = 1; i < (2 * (numRounds + 1)); ++i){
    keyArray[i] = (keyArray[i-1] + 0x9e37);
  }
  i = 0; j = 0;
  for(k = 0; k < 6 * (numRounds + 1); ++k){
    A = ROTL16( keyArray[i] + A + B, 3);
    keyArray[i] = A;
    B = ROTL16( lArray[j] + A + B, A + B);
    lArray[j] = B;
    i++;
    j++;
    i = i % (2 * (numRounds + 1));
    j = j % 2;
  }
  /*   printf("Key array for initial key %i is\n", key);
  for(i = 0; i < (2 * (numRounds + 1)); ++i){
    printf("%2i   0x%4x\n",i,keyArray[i]);
    } */
  return keyArray;
}

unsigned32 *rc5_64_keygen(unsigned64 key){
  unsigned32 *keyArray;
  unsigned32 lArray[2];
  int i = 0;
  int j = 0;
  int k = 0;
  unsigned32 A = 0;
  unsigned32 B = 0;
  lArray[0] = key;
  lArray[1] = key >> 32;
  keyArray = (unsigned32 *) malloc(sizeof(unsigned32) * 2 * (numRounds + 1)); 
  keyArray[0] = 0xb7e55163; 
  for(i = 1; i < (2 * (numRounds + 1)); ++i){
    keyArray[i] = (keyArray[i-1] + 0x9e3779b9);
  }
  i = 0; j = 0;
  for(k = 0; k < 6 * (numRounds + 1); ++k){
    A = ROTL16( keyArray[i] + A + B, 3);
    keyArray[i] = A;
    B = ROTL16( lArray[j] + A + B, A + B);
    lArray[j] = B;
    i++;
    j++;
    i = i % (2 * (numRounds + 1));
    j = j % 2;
  }
  /*   printf("Key array for initial key %i is\n", key);
  for(i = 0; i < (2 * (numRounds + 1)); ++i){
    printf("%2i   0x%4x\n",i,keyArray[i]);
    } */
  return keyArray;
}


unsigned64 rc5_64_encrypt(unsigned64 data, unsigned32 *key){
  unsigned32 a, b;
  unsigned64 result;
  int i;
  a = (data & 0xffffffff);
  b = ((data >> 32) & 0xffffffff);
  a += key[0];
  b += key[1];
  for(i = 1; i <= numRounds; ++i){
    a = a ^ b;
    a = ROTL32(a,b);
    a = a + key[2 * i];
    b = ROTL32((b ^ a), a);
    b = b + key[2 * i + 1];
  }
  result = b;
  result = result << 32;
  result = result | a;
  return result;
}

unsigned64 rc5_64_decrypt(unsigned64 data, unsigned32 *key){
  unsigned32 a, b;
  unsigned64 result;
  int i;
  a = (data & 0xffffffff);
  b = ((data >> 32) & 0xffffffff);
  for(i = numRounds; i >= 1; --i){
    b = b - key [2 * i + 1];
    b = ROTR32(b,a);
    b = b ^ a;

    a = a - key [2 * i];
    a = ROTR32(a,b);
    a = a ^ b;
  }
  b = b - key[1];
  a = a - key[0];

  result = b;
  result = result << 32;
  result = result | a;
  return result;
}



unsigned32 rc5_encrypt(unsigned32 data, unsigned16 *key){
  unsigned16 a, b;
  unsigned32 result;
  int i;
  a = (data & 0xffff);
  b = ((data >> 16) & 0xffff);
  a += key[0];
  b += key[1];
  for(i = 1; i <= numRounds; ++i){
    a = a ^ b;
    a = ROTL16(a,b);
    a = a + key[2 * i];
    b = ROTL16((b ^ a), a);
    b = b + key[2 * i + 1];
  }
  result = b;
  result = result << 16;
  result = result | a;
  return result;
}

unsigned32 rc5_decrypt(unsigned32 data, unsigned16 *key){
  unsigned16 a, b;
  unsigned32 result;
  int i;
  a = (data & 0xffff);
  b = ((data >> 16) & 0xffff);
  for(i = numRounds; i >= 1; --i){
    b = b - key [2 * i + 1];
    b = ROTR16(b,a);
    b = b ^ a;

    a = a - key [2 * i];
    a = ROTR16(a,b);
    a = a ^ b;
  }
  b = b - key[1];
  a = a - key[0];

  result = b;
  result = result << 16;
  result = result | a;
  return result;
}

typedef struct WormNode {
  unsigned32 id;
  int active;
  int infected;
  int infectedBy;
  int infectedAt; /* In ticks, not seconds! */
  unsigned32 currentlyProbing;
  /* this one uses independant 
     PRNGs for this version */
  unsigned16 *currentKey;
  int currentlyInfecting; /* A count, not an index. */
} WormNode;

int main(int argc, char **argv){
  int scanRate = 0;
  int infectTime = 0; 
  int maxInfections = 0;
  int clock = 0; /* Is in TICKS, not seconds! */
  unsigned32 scanCount = 0; /* total scans made per second */
  int hitlistSize = 0;
  int vulnerableMachines = 0;
  int activeInfections = 0;
  int totalInfections = 0;
  unsigned32 wormSeed = 0;
  unsigned32 netSeed = 1;
  unsigned16 *wormKey;
  unsigned16 *netKey;
  int startTime = 0;
  WormNode *allWorms;
  int i;

  if(sizeof(unsigned32) != 4){
    printf("Error:  unsigned32 is not 4 bytes!\n");
    exit(0);
  }
  if(sizeof(unsigned16) != 2){
    printf("Error:  unsigned16 is not 2 bytes!\n");
    exit(0);
  }
  if(sizeof(unsigned64) != 8){
    printf("Error:  unsigned64 is not 8 bytes!\n");
    exit(0);
  }
  printf("A simple simulator of a random worm's spread\n");
  if(argc != 9){
    printf("Parameters (all are required)\n");
    printf("Scans per second (int)\n");
    printf("Time to infect a machine (seconds) (int)\n");
    printf("Number of potential infections outstaning per machine (int)\n");
    printf("Size of hit list vulnerable (int)\n");
    printf("Time to process hit list (seconds) (int)\n");
    printf("Vulnerable machines in the Internet (int)\n");
    printf("Seed for worm's permutation (int)\n");
    printf("Seed for the net structure permutation (int)\n");
    exit(0);
  }
  printf("Parameters are set at\n");

  sscanf(argv[1],"%i", &scanRate);
  printf("Scans per second (%i)\n", scanRate);
  sscanf(argv[2],"%i", &infectTime);
  printf("Time to infect a machine (seconds) (%i)\n", infectTime);
  sscanf(argv[3],"%i", &maxInfections);
  printf("Number of potential infections outstaning per machine (%i)\n",
	 maxInfections);
  sscanf(argv[4],"%i", &hitlistSize);
  printf("Size of hit list vulnerable (%i)\n", hitlistSize);
  sscanf(argv[5],"%i", &startTime);
  printf("Time to process hit list (seconds) (%i)\n", startTime);
  sscanf(argv[6],"%i", &vulnerableMachines);
  printf("Vulnerable machines in the Internet (%i)\n", vulnerableMachines);
  sscanf(argv[7],"%i", &wormSeed);
  printf("Seed for worm's permutation (%i)\n", wormSeed);
  sscanf(argv[8],"%i", &netSeed);
  printf("Seed for the net structure permutation (%i)\n", netSeed);

  if(hitlistSize > vulnerableMachines){
    printf("Error, hitlist too big/total too small\n");
    exit(0);
  }
  if(startTime < 0 || startTime > 1000){
    printf("Start time out of range\n");
    exit(0);
  }
  if(wormSeed == netSeed){
    printf("Error, worm seed and net seed have to be different\n");
    exit(0);
  }


  printf("Generating worm key\n");
  wormKey = rc5_keygen(wormSeed);

  printf("Generating net key\n");
  netKey = rc5_keygen(netSeed);


  printf("Test of encryption/decryption\n");
  printf("%8x = E(0,wormKey)\n",
	 rc5_encrypt(0,wormKey));
  printf("%8x = D(E(0,wormKey),wormKey)\n",
	 rc5_decrypt(rc5_encrypt(0,wormKey),wormKey));
  printf("%8x = D(0,netKey)\n",
	 rc5_decrypt(0,netKey));
  printf("%8x = E(D(0,netKey),netKey)\n",
	 rc5_encrypt(rc5_decrypt(0,netKey),netKey));


  printf("-----\nInitializing worm world\n");
  allWorms = malloc(sizeof(WormNode) * vulnerableMachines);
  for(i = 0; i < vulnerableMachines; ++i){
    /* Set ID to be what the worm's IP address would be */
    allWorms[i].id = rc5_decrypt(i,netKey);
    allWorms[i].active = 0;
  }

  /* Initial hitlist is considered active and infected */
  for(i = 0; i < hitlistSize; ++i){
    allWorms[i].active = 1;
    allWorms[i].infected = 1;
    allWorms[i].currentlyInfecting = 0;
    
    /* Set to start infecting at the next worm in the sequence! */
    allWorms[i].currentlyProbing 
      = rc5_decrypt(allWorms[i].id,wormKey) + 1;

    /* as good a seed as I can think of right now */
    allWorms[i].currentKey = rc5_keygen(allWorms[i].id +
					   i);
    activeInfections++;
    totalInfections++;
  }


  clock = 0;
  while(totalInfections < vulnerableMachines){
    if((clock % scanRate) == 0){
      printf("At Time %3i:%02i, %7i active, %7i total, %7i scans/s\n",
	     (clock / scanRate + startTime) / 60,
	     (clock / scanRate + startTime) % 60,
	     activeInfections,
	     totalInfections, scanCount);
      fflush(stdout);
      scanCount = 0;
      if(((clock / scanRate + startTime) / 60) >= 120){
	printf("Giving up after 2 hours of simulated time\n");
	exit(0);
      }
    }
    for(i = 0; i < vulnerableMachines; ++i){
      /* this worm is alive */
      if(allWorms[i].active){
	/* changed to get target ID from PRNG local to the worm */
	unsigned32 targetId = rc5_encrypt(allWorms[i].currentlyProbing,
					  allWorms[i].currentKey);

	unsigned32 targetIndex = rc5_encrypt(targetId, netKey);

	scanCount++;
	if(targetIndex < vulnerableMachines){
	  if(allWorms[targetIndex].id != targetId){
	    printf("We haev a problem, bad net permutation!\n");
	    exit(0);
	  }
	  if(allWorms[targetIndex].active){
	    /* target machine is already active, ignore */
	    allWorms[i].currentlyProbing += 1;
	  }
	  else if(allWorms[targetIndex].infected){
	    /* Infected, but we don't know yet, just taht probe
	       failed, so move on */
	    allWorms[i].currentlyProbing += 1;
	  } else {

	    /* target is infectable, INFECT IT! */
	    allWorms[targetIndex].infected = 1;
	    allWorms[targetIndex].infectedBy = i;
	    allWorms[targetIndex].infectedAt = clock;
	    allWorms[i].currentlyInfecting += 1;

	    /* these child worms, since not a subnet scan, start at a "random"
	       place based on IP address*/
	    allWorms[targetIndex].currentlyProbing 
	      = rc5_encrypt(targetId, wormKey);
	    allWorms[targetIndex].currentKey 
	      = rc5_keygen(allWorms[targetIndex].id + targetIndex);
	    totalInfections++;
	    
	  }
	  
	} else {
	  /* target out of range/not vulnerable, increment counter */
	  allWorms[i].currentlyProbing += 1;
	}

      }
      else if(allWorms[i].infected){
	if((allWorms[i].infectedAt + (infectTime * scanRate)) < clock){
	  /* worm is now ACTIVE!  PHEAR! */
	  activeInfections++;
	  allWorms[i].active = 1;
	  allWorms[allWorms[i].infectedBy].currentlyInfecting =
	    allWorms[allWorms[i].infectedBy].currentlyInfecting - 1;
	}
      }
    }
    

    clock++;
  }

  printf("At Time %3i:%02i, total infection of %7i was achieved!\n",
	 (clock / scanRate + startTime) / 60,
	 (clock / scanRate + startTime) % 60,
	 totalInfections);
 
  return(0);

}