/* Copyright (C) 2000-2018 Peter Selinger.  
   This file is part of ccrypt. It is free software and it is covered
   by the GNU general public license. See the file COPYING for details. */

/* Try to guess ccrypt keys by exhaustively searching keys that are
   "similar" to a given key. This may be useful for recovering a
   mistyped or partially forgotten password. */

#ifdef HAVE_CONFIG_H
#include <config.h>  /* generated by configure */
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <getopt.h>

#include "lists.h"
#include "rijndael.h"
#include "readkey.h"

#define MAGIC "c051"   /* magic string for this version of ccrypt */
#define NAME "ccguess" /* name of this program */
#define CLEARLINE "\x1B[G\x1B[K"

unsigned long long int global_count = 0;

/* ---------------------------------------------------------------------- */
/* general auxiliary */

/* count number of occurrences of c in s */
int strcount(char *s, int c) {
  int count;
  
  count = 0;
  while (*s != 0) {
    if (*s == c) {
      count++;
    }
    s++;
  }
  return count;
}

/* ---------------------------------------------------------------------- */
/* from ccryptlib.c */

/* hash a keystring into a 256-bit cryptographic random value. */
static void hashstring(char *keystring, xword32 hash[8]) {
  int i;
  roundkey rkk;
  xword32 key[8];      /* rijndael key */

  for (i=0; i<8; i++) 
    key[i] = hash[i] = 0;
  
  do {
    for (i=0; i<32; i++) {
      if (*keystring != 0) {
	((xword8 *)key)[i] ^= *keystring;
	keystring++;
      }
    }
    xrijndaelKeySched(key, 256, 256, &rkk);
    xrijndaelEncrypt(hash, &rkk);
  } while (*keystring != 0);
}

/* ---------------------------------------------------------------------- */
/* key testing */

/* try the key on the given header. Return 1 if it matches, else 0. */
int try_key(roundkey *rkk, xword32 header[8]) {
  xword32 headercopy[8];
  
  memcpy(headercopy, header, 32);
  xrijndaelDecrypt(headercopy, rkk);
  if (strncmp((char *)headercopy, MAGIC, 4) != 0) {
    return 0;
  } else {
    return 1;
  }
}

/* ---------------------------------------------------------------------- */
/* patterns */

/* instead of enumerating all individual variants of a key, we enumerate
   patterns of variants. This makes it easy to weed out most duplicates. */

struct pattern_s {
  char *s; /* pattern */
  int l;   /* length */
  int w;   /* number of wildcards */
  int m;   /* number of modifications (e.g., insertions, transpositions) */
  int p;   /* negative log of probability */
  struct pattern_s *next, *next1;
};
typedef struct pattern_s pattern_t;

/* return a new pattern, or NULL with errno set on error */
pattern_t *pattern_new(char *s, int wildcard, int m, int p) {
  pattern_t *pat;

  pat = (pattern_t *)malloc(sizeof(pattern_t));
  if (pat == NULL) {
    return NULL;
  }
  pat->l = strlen(s);
  pat->s = strdup(s);
  pat->w = strcount(s, wildcard);
  pat->m = m;
  pat->p = p;
  pat->next = NULL;

  if (pat->s == NULL) {
    free(pat);
    return NULL;
  }
  return pat;
}

/* free a pattern */
void pattern_free(pattern_t *pat) {
  if (pat == NULL) {
    return;
  }
  free(pat->s);
  pat->s = NULL;
  pat->next = NULL;
  free(pat);
  return;
}

/* ---------------------------------------------------------------------- */
/* pattern generation */

/* the functions insertion, replacement, transposition, and deletion
   each generate a list of patterns without worrying about duplicates.
   They return 0 on success, or 1 on error with errno set (typically
   out of memory). */

/* prepend to plistp all patterns obtained from pat by an insertion. */
int insertion(pattern_t *pat, pattern_t **plistp, int wildcard) {
  pattern_t *plist = *plistp;
  int len = pat->l;
  char *s = pat->s;
  char buf[len+2];
  int i;
  pattern_t *pat1;
 
  for (i=0; i<=len; i++) {
    strncpy(buf, s, i);
    strcpy(buf+i+1, s+i);
    buf[i] = wildcard;
    pat1 = pattern_new(buf, wildcard, pat->m + 1, pat->p + 2);
    if (!pat1) {
      return 1;
    }
    list_prepend(plist, pat1);
  }
  *plistp = plist;
  return 0;
}

/* prepend to plistp all patterns obtained from pat by a replacement */
int replacement(pattern_t *pat, pattern_t **plistp, int wildcard) {
  pattern_t *plist = *plistp;
  int len = pat->l;
  char *s = pat->s;
  char buf[len+1];
  int i;
  pattern_t *pat1;
 
  for (i=0; i<len; i++) {
    strcpy(buf, s);
    buf[i] = wildcard;
    pat1 = pattern_new(buf, wildcard, pat->m + 1, pat->p + 1);
    if (!pat1) {
      return 1;
    }
    list_prepend(plist, pat1);
  }
  *plistp = plist;
  return 0;
}

/* prepend to plistp all patterns obtained from pat by a transposition. */
int transposition(pattern_t *pat, pattern_t **plistp, int wildcard) {
  pattern_t *plist = *plistp;
  int len = pat->l;
  char *s = pat->s;
  char buf[len+1];
  int i;
  pattern_t *pat1;
 
  for (i=0; i<len-1; i++) {
    strcpy(buf, s);
    buf[i] = s[i+1];
    buf[i+1] = s[i];
    pat1 = pattern_new(buf, wildcard, pat->m + 1, pat->p + 1);
    if (!pat1) {
      return 1;
    }
    list_prepend(plist, pat1);
  }
  *plistp = plist;
  return 0;
}

/* prepend to plistp all patterns obtained from pat by a deletion */
int deletion(pattern_t *pat, pattern_t **plistp, int wildcard) {
  pattern_t *plist = *plistp;
  int len = pat->l;
  char *s = pat->s;
  char buf[len];
  int i;
  pattern_t *pat1;
 
  for (i=0; i<len; i++) {
    strncpy(buf, s, i);
    strcpy(buf+i, s+i+1);
    pat1 = pattern_new(buf, wildcard, pat->m + 1, pat->p + 1);
    if (!pat1) {
      return 1;
    }
    list_prepend(plist, pat1);
  }
  *plistp = plist;
  return 0;
}

/* add to plistp all patterns that can be obtained from plistp by a
   single transformation; do not generate duplicates. Assume that
   plistp is sorted alphabetically, and return plistp sorted
   alphabetically. Return 0 on success, or 1 on error with errno set
   (typically out of memory). */
int expand(pattern_t **plistp, int wildcard) {
  pattern_t *rawlist = NULL;
  pattern_t *elt, *a, *b;
  pattern_t *plist = *plistp;

  /* generate single-transformation patterns */
  list_forall(elt, plist) {
    if (deletion(elt, &rawlist, wildcard) != 0) {
      return 1;
    }
    if (insertion(elt, &rawlist, wildcard) != 0) {
      return 1;
    }
    if (replacement(elt, &rawlist, wildcard) != 0) {
      return 1;
    }
    if (transposition(elt, &rawlist, wildcard) != 0) {
      return 1;
    }
  }

  /* sort them alphabetically */
  list_mergesort(pattern_t, rawlist, a, b, strcmp(a->s, b->s) <= 0);

  /* merge sorted lists plist and rawlist */
  list_merge_sorted(pattern_t, plist, rawlist, a, b, strcmp(a->s, b->s) <= 0);
  
  /* eliminate duplicates */
  list_uniq_sorted(pattern_t, plist, a, b, strcmp(a->s, b->s) == 0, a->m <= b->m, pattern_free);

  *plistp = plist;
  return 0;
}

/* list of printable characters */
char printable[] = { 
  32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48,
  49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65,
  66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82,
  83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99,
  100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112,
  113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
  126,
};

#define printable_size ((int)(sizeof(printable)/sizeof(char)))

char wildcard_chars[] = "*_#.@&$!%^-+=:";

/* Try all patterns in the list. Print out any that match. If cont is
   set, continue trying after the first match is found. Return the
   number of matching keys found. Chartable is an array of characters
   to use of size ctsize. */
int tryall(pattern_t *plist, int wildcard, char *chartable, int ctsize, xword32 (*headers)[8], int n, int cont) {
  pattern_t *pat;
  pattern_t *elt;
  pattern_t *matchlist = NULL;  /* do avoid duplicates */
  int i, j;
  int matches;
  int count = 0;
  xword32 keyblock[8];
  roundkey rkk;

  list_forall(pat, plist) {
    int w = pat->w;
    int index[w];
    int alpha[w];
    char *s = strdup(pat->s);
    int l = pat->l;
        
    fprintf(stderr, "" CLEARLINE "%s %llu", s, global_count);
    fflush(stderr);

    /* index the wildcard positions */
    i = 0;
    for (j=0; j<l; j++) {
      if (pat->s[j] == wildcard) {
	index[i] = j;
	i++;
      }
    }
    
    /* iterate through all wildcard substitutions. */
    for (i=0; i<w; i++) {
      alpha[i] = 0;
      s[index[i]] = chartable[0];
    }
    while (1) {
      global_count++;
      matches = 0;
      for (i=0; i<n; i++) {
	hashstring(s, keyblock);
	xrijndaelKeySched(keyblock, 256, 256, &rkk);
	if (try_key(&rkk, headers[i])) {
	  matches++;
	}
      }
      if (matches > 0) {
	list_find(elt, matchlist, strcmp(elt->s, s)==0);
	if (!elt) {
	  count++;
	  fprintf(stderr, "" CLEARLINE "");
	  fflush(stderr);
	  printf("\n");
	  printf("Possible match: %s (%d change%s, found after trying %llu key%s)\n", s, pat->m, pat->m==1?"":"s", global_count, global_count==1?"":"s");
	  if (matches < n) {
	    printf("Warning: key only matches %d of %d files.\n", matches, n);
	  }
	  fflush(stdout);
	  elt = pattern_new(s, wildcard, pat->m, pat->p);
	  list_append(pattern_t, matchlist, elt);
	  if (matches == n && !cont) {
	    return count;
	  }
	}
      }
      
      i = w-1;
      while (i>=0 && alpha[i] >= ctsize-1) {
        alpha[i]=0;
        s[index[i]] = chartable[0];
	i--;
      }
      if (i<0) {
	break;
      }
      alpha[i]++;
      s[index[i]] = chartable[alpha[i]];
    }
    free(s);
  }
  fprintf(stderr, "" CLEARLINE "");
  fflush(stderr);
  list_forall_unlink(elt, matchlist) {
    pattern_free(elt);
  }
  return count;
}

/* guess the keys for decrypting headers, using keys similar to "key"
   and with at most depth changes. If cont is set, continue trying
   more keys after the first match is found. Return the number of
   matches found, or -1 on failure (such as, out of memory), with
   errno set. The headers argument is an array of n headers. Chartable
   is an array of ctsize characters to use as password characters. */
int ccguess(char *key, xword32 (*headers)[8], int n, char *chartable, int ctsize, int depth, int cont) {
  int wildcard;
  pattern_t *plist;
  pattern_t *a, *b;
  int i;
  char *p;

  wildcard = '*';
  /* try to pick a wildcard character that is not part of the key. */
  for (p=wildcard_chars; *p != 0; p++) {
    if (strchr(key, *p) == NULL) {
      wildcard = *p;
      break;
    }
  }

  fprintf(stderr, "Generating patterns...");
  fflush(stderr);
  plist = pattern_new(key, wildcard, 0, 0);

  for (i=0; i<depth; i++) {
    fprintf(stderr, "%d..", i+1);
    fflush(stderr);
    if (expand(&plist, wildcard) != 0) {
      return -1;
    }
  }
  fprintf(stderr, "sorting...");
  fflush(stderr);
  /* We assume a simple statistical model where each modification
     occurs independently with probability 0.1, except insertions
     occur with probability 0.01. Therefore, the probability of a
     particular pattern with m modifications is 0.1^p. However, there
     are about 100 printable characters, so the probability of a
     particular *instance* of a pattern with m modifications and w
     wildcards is 0.1^p * 0.01^w = 0.1^(p+2w).  To minimize the
     expected time of search, we sort patterns by decreasing
     probability of their instances. */
  list_mergesort(pattern_t, plist, a, b, a->p + 2*a->w <= b->p + 2*b->w);
  fprintf(stderr, "done.\n");

  return tryall(plist, wildcard, chartable, ctsize, headers, n, cont);
}

/* ---------------------------------------------------------------------- */
/* user interface */

/* print usage information */
static void usage(FILE *fout) {
  fprintf(fout, "" NAME " " VERSION ". Search for ccrypt encryption keys.\n");
  fprintf(fout, "\n");
  fprintf(fout,
  "Usage: " NAME " [options] file...\n"
  "Options:\n"
  "    -h, --help              print this help message and exit\n"
  "    -V, --version           print version info and exit\n"
  "    -L, --license           print license info and exit\n"
  "    -K, --key <key>         specify approximate key\n"
  "    -d, --depth             try up to this many changes to key (default: 5)\n"
  "    -c, --continue          keep trying more keys after first match\n"
  "    -n, --non-printable     allow non-printable characters in keys\n"
  "    -t, --chartable <chars> characters to use in passwords (default: printable)\n"
  "Arguments:\n"
  "    file...               files to read encrypted data from, or '-' for stdin\n"
  );
}

/* print version and copyright information */
static void version(FILE *fout) {
  fprintf(fout, "" NAME " " VERSION ". Search for ccrypt encryption keys.\n");
  fprintf(fout, "Copyright (C) 2000-2018 Peter Selinger.\n");
}

static void license(FILE *fout) {
  fprintf(fout, "" NAME " " VERSION ". Search for ccrypt encryption keys.\n");
  fprintf(fout, "Copyright (C) 2000-2018 Peter Selinger.\n");
  fprintf(fout, "\n");
  fprintf(fout,
  "For the full text of the GNU General Public License, see the file\n"
  "COPYING distributed with this software.\n"
  "\n"
  "This program is free software; you can redistribute it and/or modify\n"
  "it under the terms of the GNU General Public License as published by\n"
  "the Free Software Foundation; either version 2 of the License, or\n"
  "(at your option) any later version.\n"
  "\n"
  "This program is distributed in the hope that it will be useful,\n"
  "but WITHOUT ANY WARRANTY; without even the implied warranty of\n"
  "MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n"
  "GNU General Public License for more details.\n"
  "\n"
  "You should have received a copy of the GNU General Public License\n"
  "along with this program; if not, write to the Free Software Foundation,\n"
  "Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.\n"
  );
}

static struct option longopts[] = {
  {"help",          0, 0, 'h'},
  {"version",       0, 0, 'V'},
  {"license",       0, 0, 'L'},
  {"key",           1, 0, 'K'},
  {"depth",         1, 0, 'd'},
  {"continue",      0, 0, 'c'},
  {"non-printable", 0, 0, 'n'},
  {"chartable",     1, 0, 't'},
  {0, 0, 0, 0}
};

static const char *shortopts = "hVLK:d:cnt:";

/* structure to hold command-line */
typedef struct {
  char *keyword;
  int depth;
  char **infiles;    /* input filenames ("-" for stdin) */
  int nfiles;        /* number of filenames */
  int cont;
  int non_printable;
  char *chartable;
} cmdline;

static cmdline read_commandline(int ac, char *av[]) {
  cmdline cmd;
  int c;
  char *p;

  /* defaults: */
  cmd.keyword = NULL;
  cmd.depth = 5;
  cmd.nfiles = 0;
  cmd.cont = 0;
  cmd.non_printable = 0;
  cmd.chartable = NULL;

  while ((c = getopt_long(ac, av, shortopts, longopts, NULL)) != -1) {
    switch (c) {
    case 'h':
      usage(stdout);
      exit(0);
      break;
    case 'V':
      version(stdout);
      exit(0);
      break;
    case 'L':
      license(stdout);
      exit(0);
      break;
    case 'K':
      free(cmd.keyword);
      cmd.keyword = strdup(optarg);
      if (!cmd.keyword) {
	goto mem_error;
      }
      /* attempt to erase keyword from command line so that subsequent
         calls to 'ps' don't display it */
      for (p=optarg; *p; p++) {  
	*p = 0;
      }
      break;
    case 'd':
      cmd.depth = strtod(optarg, &p);
      if (*p || cmd.depth <= 0) {
        fprintf(stderr, "" NAME ": invalid depth -- %s\n", optarg);
        exit(2);
      }
      break;
    case 'c':
      cmd.cont = 1;
      break;
    case 'n':
      cmd.non_printable = 1;
      cmd.chartable = NULL;
      break;
    case 't':
      cmd.chartable = optarg;
      cmd.non_printable = 0;
      break;
    case '?':
      fprintf(stderr, "Try --help for more information.\n");
      exit(2);
      break;
    default:
      fprintf(stderr, "" NAME ": unimplemented option -- %c\n", c);
      exit(2);
    }
  }

  switch (ac-optind) {
  case 0:
    fprintf(stderr, "" NAME ": missing filename. Try --help for more information.\n");
    exit(2);
    break;
  default:
    cmd.infiles = &av[optind];
    cmd.nfiles = ac-optind;
    break;
  }
  return cmd;

 mem_error:
  fprintf(stderr, "" NAME ": %s\n", strerror(errno));
  free(cmd.keyword);
  exit(2);
}

int main(int ac, char **av) {
  xword32 (*headers)[8];
  int r;
  cmdline cmd;
  FILE *fin;
  int i, j;
  char *file;
  char bytes[256];
  char *chartable;
  int ctsize;

  /* read command line */
  cmd = read_commandline(ac, av);

  /* prepare character table */
  if (cmd.chartable) {
    /* condense character table to eliminate duplicates */
    memset(bytes, 0, 256);
    for (i=0; cmd.chartable[i]; i++) {
      bytes[((int)cmd.chartable[i]) & 0xff] = cmd.chartable[i];
    }
    for (i=0, j=0; i<256; i++) {
      if (bytes[i]) {
	bytes[j] = bytes[i];
	j++;
      }
    }
    chartable = bytes;
    ctsize = j;
  } else if (cmd.non_printable) {
    for (i=1; i<256; i++) {
      bytes[i-1] = i;
    }
    chartable = bytes;
    ctsize = 255;
  } else {
    chartable = printable;
    ctsize = printable_size;
  }

  /* read keyword from terminal if necessary */
  if (cmd.keyword==NULL) {
    cmd.keyword = readkey("Enter approximate key: ", "",  NAME , 1);
    if (cmd.keyword==NULL) {  /* end of file: exit gracefully */
      fprintf(stderr, "" NAME ": no key given\n");
      return 2;
    }
  }

  headers = (xword32 (*)[8]) malloc(cmd.nfiles * sizeof(xword32[8]));
  if (!headers) {
    fprintf(stderr, "" NAME ": %s\n", strerror(errno));
    return 2;
  }

  /* read headers from all files */
  for (i=0; i<cmd.nfiles; i++) {
    file = cmd.infiles[i];

    /* open encrypted file */
    if (strcmp(file, "-") == 0) {
      fin = stdin;
    } else {
      fin = fopen(file, "r");
      if (!fin) {
	fprintf(stderr, "" NAME ": could not read %s: %s\n", file, strerror(errno));
	return 2;
      }
    }
    
    /* read header */
    r = fread(headers[i], 32, 1, fin);
    if (r != 1) {
      fprintf(stderr, "" NAME ": %s: file too short\n", file);
      return 2;
    }
    
    /* close file */
    if (strcmp(file, "-") != 0) {
      fclose(fin);
    }
  }

  /* payload */
  r = ccguess(cmd.keyword, headers, cmd.nfiles, chartable, ctsize, cmd.depth, cmd.cont);
  free(headers);

  if (r < 0) {
    fprintf(stderr, "" NAME ": %s\n", strerror(errno));
    return 2;
  } 
  return r>0 ? 0 : 1;
}
