/*
discrete log cracker
breaks a discrete log problem via pollard's rho algorithm

written by seth hardy, shardy@aculei.net

license-ish thing:
you can use this code if you give credit where credit is due.
I assume no responsibility for what happens if you use it, though.

*/  


#include <stdio.h>
#include <time.h>
#include "gmp.h"

#define breakpoint 10000
#define MAXSIZE 100
#define FALSE 0
#define TRUE  1

void about(void);

main(int argc,char *argv[]) {

  mpz_t a,b,p,n;
  mpz_t xi,ai,bi,x2i,a2i,b2i;
  mpz_t temp,temp2,count;
  
  int debug;
  int timing;
  int testing;

  unsigned long start,end;

  int loop;
  size_t size;

  char instr[MAXSIZE];

  debug = timing = FALSE;

  if (argc>1) { 
    if((strcmp(argv[1],"-h") == 0) || (strcmp(argv[1],"-H") == 0))
      about();
    for(testing=1;testing<argc;testing++) {
      if ((strcmp(argv[testing],"-d") == 0) || (strcmp(argv[testing],"-D")
        == 0)) debug = TRUE;
      if ((strcmp(argv[testing],"-t") == 0) || (strcmp(argv[testing],"-T")
        == 0)) timing = TRUE;
    }
  }
  

  srand(time(NULL));

  mpz_init(a);
  mpz_init(b);
  mpz_init(p);
  mpz_init(n);

  mpz_init(xi);
  mpz_init(ai);
  mpz_init(bi);
  mpz_init(x2i);
  mpz_init(a2i);
  mpz_init(b2i);
  mpz_init(temp);
  mpz_init(temp2);
  mpz_init(count);

  printf("\ndlp-crack starting...\n\n");
  printf("          Alpha: ");
  scanf("%s",instr);
  if(debug) printf("%s scanned as Alpha.\n",instr);
  mpz_set_str(a,instr,10);
  printf("           Beta: ");
  scanf("%s",instr);
  if(debug) printf("%s scanned as Beta.\n",instr);
  mpz_set_str(b,instr,10);
  printf("Prime modulus p: ");
  scanf("%s",instr);
  if(debug) printf("%s scanned as p.\n",instr);
  mpz_set_str(p,instr,10);
  printf("  Group order n: ");
  scanf("%s",instr);
  if(debug) printf("%s scanned as n.\n",instr);
  mpz_set_str(n,instr,10);
  printf("\n");

  // Begin cracking!

  start = clock();
  size = mpz_size(n);
  testing = 1;

  mpz_set_ui(count,0);
  mpz_random(ai,size-1);
  mpz_random(bi,size-1);

  mpz_powm(temp,a,ai,p);
  mpz_powm(xi,b,bi,p);
  mpz_mul(xi,xi,temp);
  mpz_mod(xi,xi,p);

  mpz_set(x2i,xi);
  mpz_set(a2i,ai);
  mpz_set(b2i,bi);

  do {
    if(debug) {
       printf("i = ");
       mpz_out_str(stdout,10,count);
       printf(": ");
       printf("(xi,ai,bi) = ");
       mpz_out_str(stdout,10,xi);
       printf(",");
       mpz_out_str(stdout,10,ai);
       printf(",");
       mpz_out_str(stdout,10,bi);
       printf(" -- (x2i,a2i,b2i) = ");
       mpz_out_str(stdout,10,x2i);
       printf(",");
       mpz_out_str(stdout,10,a2i);
       printf(",");
       mpz_out_str(stdout,10,b2i);
       printf("\n");

       if(testing > breakpoint) {
         printf("Program breakpoint hit. Terminating...\n");
         return;
       }
  }

  mpz_add_ui(count,count,1);

  switch(mpz_mod_ui(temp,xi,3))
  {
    case 0 : mpz_mul(xi,xi,xi);
             mpz_mul_ui(ai,ai,2);
             mpz_mul_ui(bi,bi,2);
             break;
    case 1 : mpz_mul(xi,xi,b);
             mpz_add_ui(bi,bi,1);
             break;
    case 2 : mpz_mul(xi,xi,a);
             mpz_add_ui(ai,ai,1);
             break;
  }

  mpz_mod(xi,xi,p);
  mpz_mod(ai,ai,n);
  mpz_mod(bi,bi,n);

  switch(mpz_mod_ui(temp,x2i,3))
  {
    case 0 : mpz_mul(x2i,x2i,x2i);
             mpz_mul_ui(a2i,a2i,2);
             mpz_mul_ui(b2i,b2i,2);
             break;
    case 1 : mpz_mul(x2i,x2i,b);
             mpz_add_ui(b2i,b2i,1);
             break;
    case 2 : mpz_mul(x2i,x2i,a);
             mpz_add_ui(a2i,a2i,1);
             break;
  }

  mpz_mod(x2i,x2i,p);
		
  switch(mpz_mod_ui(temp,x2i,3))
  {
    case 0 : mpz_mul(x2i,x2i,x2i);
             mpz_mul_ui(a2i,a2i,2);
             mpz_mul_ui(b2i,b2i,2);
             break;
    case 1 : mpz_mul(x2i,x2i,b);
             mpz_add_ui(b2i,b2i,1);
             break;
    case 2 : mpz_mul(x2i,x2i,a);
             mpz_add_ui(a2i,a2i,1);
             break;
  }

  mpz_mod(x2i,x2i,p);
  mpz_mod(a2i,a2i,n);
  mpz_mod(b2i,b2i,n);

  if(debug) testing++;	

  } while(mpz_cmp(xi,x2i) != 0);
	
  if(debug) {
    printf("Collision found!\n");
    printf("xi = ");
    mpz_out_str(stdout,10,xi);
    printf("   x2i = ");
    mpz_out_str(stdout,10,x2i);
    printf("\nai = ");
    mpz_out_str(stdout,10,ai);
    printf("   a2i = ");
    mpz_out_str(stdout,10,a2i);
    printf("\nbi = ");
    mpz_out_str(stdout,10,bi);
    printf("   b2i = ");
    mpz_out_str(stdout,10,b2i);
    printf("\n");
    }
	
  mpz_sub(temp,bi,b2i);
  mpz_mod(temp,temp,n);
	
  if(mpz_cmp_ui(temp,0) == 0) printf("Solution failed.\n");
  else {
    mpz_sub(temp2,a2i,ai);
    mpz_invert(temp,temp,n);
    mpz_mul(temp,temp,temp2);
    mpz_mod(temp,temp,n);
  }

  end = clock();

  printf("Solution for DLP: ");
  mpz_out_str(stdout,10,temp);
  printf("\nIterations: ");
  mpz_out_str(stdout,10,count);
  printf("\n");
  if(timing) 
    printf("Time taken (sec): %f\n",(float)(end-start)/1000000);

  mpz_clear(a);
  mpz_clear(b);
  mpz_clear(n);
  mpz_clear(p);

  mpz_clear(xi);
  mpz_clear(x2i);
  mpz_clear(ai);
  mpz_clear(a2i);
  mpz_clear(bi);
  mpz_clear(b2i);
  mpz_clear(temp);
  mpz_clear(count);
  mpz_clear(temp2);

  printf("dlp-crack terminated successfully.\n\n");

  return;

}

void about(void) {

  printf("\n\ndlp-crack v1.0 -- Discrete Logarithm Problem Cracker\n\n");
  printf("USAGE: dlp-crack [options]\n\n");
  printf("\t-d   Show debug information\n");
  printf("\t-t   Show timing information\n\n\n");
  exit(0);

}