/**
 * Copyright (C) 2006 Henning Norén
 * Copyright (C) 1996-2005 Glyph & Cog, LLC.
 * 
 * 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., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, 
 * USA.
 */

#include <stdio.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <ctype.h>
#include "pdfcrack.h"
#include "md5.h"
#include "rc4.h"
#include "passwords.h"

static const uint8_t
pad[32] = {
  0x28,0xBF,0x4E,0x5E,0x4E,0x75,0x8A,0x41,
  0x64,0x00,0x4E,0x56,0xFF,0xFA,0x01,0x08,
  0x2E,0x2E,0x00,0xB6,0xD0,0x68,0x3E,0x80,
  0x2F,0x0C,0xA9,0xFE,0x64,0x53,0x69,0x7A
};

/**
static void
padPassword(uint8_t *padpwd, const uint8_t *pwd, const unsigned int pwd_l) {
  if(pwd_l < 32) {
    memcpy(padpwd, pwd, pwd_l);
    memcpy(padpwd + pwd_l, pad, 32-pwd_l);
  }
  else
    memcpy(padpwd, pwd, 32);
}
*/

static uint8_t *encKeyWorkSpace;
static uint8_t password_user[32];
static const uint8_t *rev3TestKey;
static unsigned int ekwlen;
static unsigned int nrprocessed;
static unsigned int currPWLen;
static clock_t startTime;
static const EncData *encdata;
static bool crackDone;
static bool knownPassword;
static bool workWithUser;
static uint8_t *currPW;


bool
printProgress(void) {
  clock_t currentTime;
  char str[33];

  if(crackDone)
    return true;

  currentTime = clock();
  memcpy(str,currPW,currPWLen);
  str[currPWLen] = '\0';
  printf("Average Speed: %.1f w/s. ",
	 nrprocessed/(((double)(currentTime-startTime))/CLOCKS_PER_SEC));
  printf("Current Word: '%s'\n",str);
  nrprocessed = 0;
  startTime = clock();
  return false;
}

/**
 * Initialisation of the encryption key workspace to manage a bit faster 
 * switching between keys
 */
static unsigned int
initEncKeyWorkSpace(const int revision, const bool encMetaData, 
		    const int permissions, const uint8_t *ownerkey, 
		    const uint8_t *fileID, const unsigned int fileIDLen) {
  /** 
   *   Algorithm 3.2 Computing an encryption key (PDF Reference, v 1.6, p.101)
   *
   *   Make space for:
   *   field           | bits
   *   ----------------------
   *   padded password | 32
   *   O entry         | 32
   *   P entry         |  4
   *   fileID          | <fileIDLEn>
   *  [extra padding]  | [4] (Special for step 6)
   **/
  unsigned int size = (revision > 2 && !encMetaData) ? 72 : 68;
  encKeyWorkSpace = malloc(size + fileIDLen);

  /** Just to be sure we have no uninitalized stuff in the workspace */
  memcpy(encKeyWorkSpace, pad, 32);

  /** 3 */ 
  memcpy(encKeyWorkSpace + 32, ownerkey, 32);

  /** 4 */
  encKeyWorkSpace[64] = permissions & 0xff;
  encKeyWorkSpace[65] = (permissions >> 8) & 0xff;
  encKeyWorkSpace[66] = (permissions >> 16) & 0xff;
  encKeyWorkSpace[67] = (permissions >> 24) & 0xff;

  /** 5 */
  memcpy(encKeyWorkSpace + 68, fileID, fileIDLen);

  /** 6 */
  if(revision > 2 && !encMetaData) {
    encKeyWorkSpace[68+fileIDLen] = 0xff;
    encKeyWorkSpace[69+fileIDLen] = 0xff;
    encKeyWorkSpace[70+fileIDLen] = 0xff;
    encKeyWorkSpace[71+fileIDLen] = 0xff;
  }

  return size+fileIDLen;
}

/**
static void
printString(const uint8_t *str, const unsigned int len) {
  unsigned int i;
  for(i=0;i<len;i++)
    printf("%x ",str[i]);
  printf("\n");
}
*/

static uint8_t
isolat1ToUpper(const uint8_t b) {
  if(unlikely(b >= 0xe0 && b <= 0xf6))
    return b-0x20;
  else
    return toupper(b);
}

static bool
permutate(void) {
  static bool ret = true;
  
  currPW[0] = isolat1ToUpper(currPW[0]);
  ret = !ret;
  
  return ret;
}

static void
foundPassword(void) {
  char str[33];

  memcpy(str,currPW,currPWLen);
  str[currPWLen] = '\0';
  printf("found %s-password: '%s'\n", workWithUser?"user":"owner", str);
}

static bool
isUserPasswordRev2(void) {
  uint8_t enckey[16];

  md5(encKeyWorkSpace, ekwlen, enckey);

  return rc4Match40b(enckey, encdata->u_string, 32, pad);
}

static bool
isUserPasswordRev3(void) {
  uint8_t test[16], enckey[16], tmpkey[16];
  int i;
  unsigned int length, j;

  length = encdata->length/8;
  md5(encKeyWorkSpace, ekwlen, enckey);      
  md5_50(enckey);
  memcpy(test, encdata->u_string, 16);

  for(i = 19; i >= 0; --i) {
    for(j = 0; j < length; ++j)
      tmpkey[j] = enckey[j] ^ i;
    
    rc4Decrypt(tmpkey, test, 3, test);
  }
      
  /** if partial test succeeds we make a full check to make sure */
  if(unlikely(memcmp(test, rev3TestKey, 3) == 0)) {
    memcpy(test, encdata->u_string, 16);
    for(i = 19; i >= 0; --i) {
      for(j = 0; j < length; ++j)
	tmpkey[j] = enckey[j] ^ i;

      rc4Decrypt(tmpkey, test, 16, test);
    }
    if(memcmp(test, rev3TestKey, 16) == 0) {
      return true;
    }
  }
  return false;
}

#define BEGIN_CRACK_LOOP() {				\
    currPWLen = setPassword(currPW);			\
    if(unlikely(lpasslength != currPWLen)) {		\
      if(likely(currPWLen < 32))			\
	memcpy(currPW + currPWLen, pad, 32-currPWLen);	\
      lpasslength = currPWLen;				\
    }							\
  }

static bool
runCrackRev2_o(void) {
  uint8_t enckey[16];
  unsigned int lpasslength;

  lpasslength = 0;
  startTime = clock();
  while(nextPassword()) {
    BEGIN_CRACK_LOOP();
  
    do {
      md5(currPW, 32, enckey);

      rc4Decrypt(enckey, encdata->o_string, 32, encKeyWorkSpace);
      if(isUserPasswordRev2())
	return true;

      ++nrprocessed;
    } while(permutate());
  }
  return false;
}

static bool
runCrackRev3_o(void) {
  uint8_t test[32], enckey[16], tmpkey[16];
  unsigned int j, length, lpasslength;
  int i;

  length = encdata->length/8;
  lpasslength = 0;
  startTime = clock();
  while(nextPassword()) {
    BEGIN_CRACK_LOOP();

    do {
      md5(currPW, 32, enckey);
      
      md5_50(enckey);

      memcpy(test, encdata->o_string, 32);
      for(i = 19; i >= 0; --i) {
	for(j = 0; j < length; ++j)
	  tmpkey[j] = enckey[j] ^ i;
	rc4Decrypt(tmpkey, test, 32, test);
      }

      memcpy(encKeyWorkSpace, test, 32);

      if(isUserPasswordRev3())
	return true;

      ++nrprocessed;
    } while(permutate());
  }
  return false;
}

static bool
runCrackRev2_of(void) {
  uint8_t enckey[16];
  unsigned int lpasslength;

  lpasslength = 0;
  startTime = clock();
  while(nextPassword()) {
    BEGIN_CRACK_LOOP();

    do {
      md5(encKeyWorkSpace, 32, enckey);

      /* Algorithm 3.4 reversed */
      if(rc4Match40b(enckey, encdata->o_string, 32, password_user))
	return true;

      ++nrprocessed;
    } while(permutate());
  }
  return false;
}

static bool
runCrackRev3_of(void) {
  uint8_t test[32], enckey[16], tmpkey[16];
  unsigned int j, length, lpasslength;
  int i;

  length = encdata->length/8;
  lpasslength = 0;
  startTime = clock();
  while(nextPassword()) {
    BEGIN_CRACK_LOOP();
    
    do {
      md5(encKeyWorkSpace, 32, enckey);
      
      md5_50(enckey);

      memcpy(test, encdata->o_string, 32);
      for(i = 19; i >= 0; --i) {
	for(j = 0; j < length; ++j)
	  tmpkey[j] = enckey[j] ^ i;
	
	rc4Decrypt(tmpkey, test, 3, test);
      }

      /** if partial test succeeds we make a full check to make sure */
      if(unlikely(memcmp(test, password_user, 3) == 0)) {
	memcpy(test, encdata->o_string, 32);
	for(i = 19; i >= 0; --i) {
	  for(j = 0; j < length; ++j)
	    tmpkey[j] = enckey[j] ^ i;
	  rc4Decrypt(tmpkey, test, 32, test);
	}

	if(memcmp(test, password_user, 32) == 0)
	  return true;
      }

      ++nrprocessed;
    } while(permutate());
  }
  return false;
}

static bool
runCrackRev3(void) {
  uint8_t test[16], enckey[16], tmpkey[16];
  unsigned int j, length, lpasslength;
  int i;

  length = encdata->length/8;
  lpasslength = 0;
  startTime = clock();
  while(nextPassword()) {
    BEGIN_CRACK_LOOP();

    do {
      md5(encKeyWorkSpace, ekwlen, enckey);

      md5_50(enckey);
      memcpy(test, encdata->u_string, 16);
      /** Algorithm 3.5 reversed 
	  Only do a partial test for the 3 first charcters in the string.
	  3 should be enough to avoid collisions too often as it only should
	  happen once every 256^3=16777216 time and that is seldom enough to
	  make a complete retest of the full string.
	  Lowering the number of characters tried to 2 is rather pointless 
	  as the decryption is cheap.
      */
      
      for(i = 19; i >= 0; --i) {
	for(j = 0; j < length; ++j)
	  tmpkey[j] = enckey[j] ^ i;

	rc4Decrypt(tmpkey, test, 3, test);
      }

      /** if partial test succeeds we make a full check to make sure */
      if(unlikely(memcmp(test, rev3TestKey, 3) == 0)) {
	memcpy(test, encdata->u_string, 16);
	for(i = 19; i >= 0; --i) {
	  for(j = 0; j < length; ++j)
	    tmpkey[j] = enckey[j] ^ i;
	  rc4Decrypt(tmpkey, test, 16, test);
	}	  
	if(memcmp(test, rev3TestKey, 16) == 0)
	  return true;
      }

      ++nrprocessed;
    } while(permutate());
  }
  return false;
}


static bool
runCrackRev2(void) {
  uint8_t enckey[16];
  unsigned int lpasslength;

  lpasslength = 0;
  startTime = clock();
  while(nextPassword()) {
    BEGIN_CRACK_LOOP();

    do {
      md5(encKeyWorkSpace, ekwlen, enckey);

      /* Algorithm 3.4 reversed */
      if(rc4Match40b(enckey, encdata->u_string, 32, pad))
	return true;

      ++nrprocessed;
    } while(permutate());
  }
  return false;
}

void
runCrack(void) {
  bool found = false;
  uint8_t cpw[32];
  if(!workWithUser && !knownPassword) {
    memcpy(cpw, pad, 32);
    currPW = cpw;
    if(encdata->revision == 2)
      found = runCrackRev2_o();
    else
      found = runCrackRev3_o();
  }  
  else if(encdata->revision == 2) {
    if(workWithUser)
      found = runCrackRev2();
    else /** knownPassword */
      found = runCrackRev2_of();
  }
  else {
    if(workWithUser)
      found = runCrackRev3();
    else /** knownPassword */
      found = runCrackRev3_of();
  }
  crackDone = true;
  if(!found)
    printf("Could not find password\n");
  else
    foundPassword();
  currPW = NULL;
}

void
cleanPDFCrack(void) {
  if(rev3TestKey)
    free(rev3TestKey);
  if(encKeyWorkSpace)
    free(encKeyWorkSpace);
}

bool
initPDFCrack(const EncData *e, const uint8_t *upw, const bool user,
	     const char *wl, const passwordMethod pm, FILE *file, 
	     const char *cs, const unsigned int minPw, 
	     const unsigned int maxPw) {
  uint8_t buf[128];
  unsigned int upwlen;
  uint8_t *tmp;

  ekwlen = initEncKeyWorkSpace(e->revision, e->encryptMetaData, e->permissions,
			       e->o_string, e->fileID, e->fileIDLen);

  encdata = e;
  currPW = encKeyWorkSpace;
  currPWLen = 0;
  nrprocessed = 0;
  workWithUser = user;
  crackDone = false;
  setrc4DecryptMethod((unsigned int)e->length);
  if(upw) {
    upwlen = strlen((char*)upw);
    if(upwlen > 32)
      upwlen = 32;
    memcpy(password_user, upw, upwlen);
    memcpy(password_user+upwlen, pad, 32-upwlen);
    memcpy(encKeyWorkSpace, password_user, 32);
    knownPassword = true;
  }

  if(encdata->revision == 2) {
    if(knownPassword) {
      if(!isUserPasswordRev2())
	return false;
      memcpy(encKeyWorkSpace, pad, 32);
    }
    else {
      memcpy(password_user, pad, 32);
      knownPassword = isUserPasswordRev2();
    }
  }
  else if(e->revision == 3) {
    memcpy(buf, pad, 32);
    memcpy(buf + 32, e->fileID, e->fileIDLen);
    tmp = malloc(sizeof(uint8_t)*16);
    md5(buf, 32+e->fileIDLen, tmp);
    rev3TestKey = tmp;
    if(knownPassword) {
      if(!isUserPasswordRev3())
	return false;
      memcpy(encKeyWorkSpace, pad, 32);
    }
    else {
      memcpy(password_user, pad, 32);
      knownPassword = isUserPasswordRev3();
    }
  }
  
  initPasswords(pm, file, wl, cs, minPw, maxPw);
  return true;
}

void
loadState(FILE *file, EncData *e, char **wl, bool *user) {
  unsigned int i;
  int tmp;

  /** possible overflow, but I really don't care */
  e->s_handler = malloc(sizeof(uint8_t)*32);
  fscanf(file,"PDF: %d.%d\nR: %d\nV: %d\nP: %d\nL: %d\n"
	 "Filter: %[^\n]\nFileIDLen: %d\nMetaData: %d\nFileID: ",
	 &e->version_major, &e->version_minor, &e->revision, &e->version,
	 &e->permissions, &e->length, e->s_handler, &e->fileIDLen, &tmp);
  e->encryptMetaData = (tmp == true);
  e->fileID = malloc(sizeof(uint8_t)*e->fileIDLen);
  e->o_string = malloc(sizeof(uint8_t)*32);
  e->u_string = malloc(sizeof(uint8_t)*32);
  for(i=0;i<e->fileIDLen;i++) {
    fscanf(file, "%d ", &tmp);
    e->fileID[i] = tmp;
  }
  fscanf(file,"\nO: ");
  for(i=0;i<32;i++) {
    fscanf(file, "%d ", &tmp);
    e->o_string[i] = tmp;
  }
  fscanf(file,"\nU: ");
  for(i=0;i<32;i++) {
    fscanf(file, "%d ", &tmp);
    e->u_string[i] = tmp;
  }
  fscanf(file, "\nUser: %d", &tmp);
  *user = (tmp == true);
  fscanf(file, "\nUserPw: %d\n", &tmp);
  knownPassword = (tmp == true);
  if(knownPassword)
    for(i=0;i<32;i++) {
      fscanf(file, "%d ", &tmp);
      password_user[i] = tmp;
    }
  fscanf(file, "\n");
  pw_loadState(file, wl);
}

void
saveState(FILE *file) {
  unsigned int i;

  fprintf(file, "PDF: %d.%d\nR: %d\nV: %d\nP: %d\nL: %d\n"
	  "Filter: %s\nFileIDLen: %d\nMetaData: %d\nFileID: ",
	  encdata->version_major, encdata->version_minor, encdata->revision,
	  encdata->version, encdata->permissions, encdata->length,
	  encdata->s_handler,encdata->fileIDLen,(int)encdata->encryptMetaData);
  for(i=0;i<encdata->fileIDLen;i++)
    fprintf(file, "%d ", encdata->fileID[i]);
  fprintf(file,"\nO: ");
  for(i=0;i<32;i++)
    fprintf(file, "%d ", encdata->o_string[i]);
  fprintf(file,"\nU: ");
  for(i=0;i<32;i++)
    fprintf(file, "%d ", encdata->u_string[i]);
  fprintf(file, "\nUser: %d\nUserPw: %d\n",
	  (int)workWithUser, (int)knownPassword);
  if(knownPassword) {
    for(i=0;i<32;i++)
      fprintf(file, "%d ", password_user[i]);
  }
  fprintf(file, "\n");
  pw_saveState(file);
}
