#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include <time.h>
#include <stdlib.h>
#include <math.h>

#ifdef __VSDSP__
#include <vstypes.h>
#else
#include <inttypes.h>
typedef uint32_t u_int32;
typedef uint16_t u_int16;
#define __mem_y
#define memcpyYX memcpy
#endif

#include "chacha20.h"



#define OFFSET     0  /* Must be positive, 0 won't do (yes it does) */
#define MULTIPLIER 1  /* Turning up to 10 may enhance compression */

struct TREE {
  struct TREE *left, *right, *next, *prev;
  int c, f, index;
};

static struct TREE *base = NULL;
static double acBitsTotal;

static void Remove(struct TREE *t);

void FreeLeftRight(struct TREE *t) {
  if (t->left) {
    FreeLeftRight(t->left);
    free(t->left);
  }
  if (t->right) {
    FreeLeftRight(t->right);
    free(t->right);
  }
}

void FreeTree(struct TREE *t) {
  if (t) {
    FreeLeftRight(t);
    FreeTree(t->next);
    free(t);
  }
}

struct TREE *AllocTN(void) {
  struct TREE *t = calloc(1, sizeof(struct TREE));
  static int index = 0;
  t->index = index++;
  if (!t) {
    fprintf(stderr, "OUT OF MEMORY\n");
    exit(EXIT_FAILURE);
  }
  return t;
}


static void Remove(struct TREE *t) {
  if (t->prev) {
    t->prev->next = t->next;
  }
  if (t->next) {
    t->next->prev = t->prev;
  }
  t->prev = NULL;
  t->next = NULL;
}


static void AddAfter(struct TREE *d, struct TREE *n) {
  n->next = d->next;
  if (n->next) {
    n->next->prev = n;
  }
  n->prev = d;
  d->next = n;
}


static void AddLeft(struct TREE *d, struct TREE *n) {
  d->left = n;
}

static void AddRight(struct TREE *d, struct TREE *n) {
  d->right = n;
}

static void PrintBin(int n, int bits) {
  int i = bits;
  while (i-- > 0) {
    printf("%c", (n & (1<<i)) ? '1' : '0');
  }
}

static int Bits(struct TREE *n, int depth, int print, int code, int chars) {
  if (!depth) {
    acBitsTotal = 0.0;
  }
  if (n) {
    if (n->c != EOF) {
      double acBits =
	log(chars/(double)((n->f-OFFSET)/MULTIPLIER))/log(2.0);
      acBitsTotal += acBits*(n->f-OFFSET)/MULTIPLIER;

      if (depth == print) {
	printf("%c (%2x), f=%5d, hBits=%2d, aBits=%5.2f, hCode=",
	       ((isprint(n->c)) ? n->c : '?'), n->c,
	       (n->f-OFFSET)/MULTIPLIER, depth, acBits);
	PrintBin(code, depth);
	printf("\n");
      }
    }
    return ((n->c != EOF) ? (n->f-OFFSET)*depth/MULTIPLIER : 0) +
      Bits(n->left,  depth+1, print, (code<<1) + 0, chars) +
      Bits(n->right, depth+1, print, (code<<1) + 1, chars);
  } else {
    return 0;
  }
}

const char *progName = NULL;

int EstimateEntropy(const unsigned char *s, int verbose) {
  int c, i, chars = 0, ch, ca;
  int tab[256];
  int aritEnt = 0, lenEnt = 0;
  
  if (verbose) {
    printf("EstimateEntropy(\"%s\")\n", s);
  }
  
  for (i=0; i<256; i++) {
    tab[i] = OFFSET;
  }

  while ((c = *s++)) {
    tab[c] += MULTIPLIER;
    chars++;
  }

  base = AllocTN();

  c = 0;
  if (verbose > 1) {
    printf("Missing chars:\n");
  }
  for (i=0; i<256; i++) {
    if (tab[i] > OFFSET) {
      struct TREE *n = AllocTN();
      n->c = i;
      n->f = tab[i];
      AddAfter(base, n);
    } else {
      if (verbose > 1) {
	printf("%c %02x%c", ((isprint(i)) ? i : '?'), i, (++c%16) ? ' ' : '\n');
      }
    }
  }
  if (c%16 && verbose > 1) {
    printf("\n");
  }

  if (verbose) {
    printf("Coded words:\n");
  }

  while (chars && base->next->next) {
    struct TREE *low1TN = base->next, *low2TN = low1TN->next,
      *n = low2TN->next;
    if (low1TN->f > low2TN->f) {
      struct TREE *tn = low1TN;
      low1TN = low2TN;
      low2TN = tn;
    }
    while (n) {
      if (n->f < low2TN->f) {
	if (low1TN->f > n->f) {
	  low2TN = low1TN;
	  low1TN = n;
	} else {
	  low2TN = n;
	}
      }
      n = n->next;
    }
    n = AllocTN();
    n->c = EOF;
    n->f = low1TN->f + low2TN->f;
    AddAfter(base, n);
    Remove(low1TN);
    Remove(low2TN);
    AddLeft(n, low1TN);
    AddRight(n, low2TN);
  }

  if (verbose) {
    for (i=32; i>0; i--) {
      Bits(base->next, 0, i, 0, chars);
    }
  }
  i = Bits(base->next, 0, -1, 0, chars);
  ch = chars ? (((i+7)/8)*200/chars+1)/2 : 100;
  ca = chars ? (((int)(ceil(acBitsTotal)+7)/8)*200/chars+1)/2 : 100;

  FreeTree(base);
  base = NULL;

  aritEnt = acBitsTotal;
  lenEnt = log(chars)/log(2);
  
  if (verbose) {
    printf("Original number of chrs %7d, bits = %7d (100%%) (  0%%)\n",
	   chars, chars*8);
    printf("Huffman compr no of crs %7d, bits = %7d (%3d%%) (%3d%%)\n",
	   (i+7)/8, i, ch, 100-ch);
    printf("Arithmetic cpr n of crs %7d, bits = %7d (%3d%%) (%3d%%)\n",
	   ((int)ceil(acBitsTotal)+7)/8, (int)(ceil(acBitsTotal)),
	   ca, 100-ca);
    printf("Total enthropy estimate: %d(arithmetic) + %d(length) = %d bits\n",
	   aritEnt, lenEnt, aritEnt+lenEnt);
  }

  return aritEnt+lenEnt;
}




unsigned ReadKey(const char *fname, u_int32 *key) {
  FILE *kf;
  char *line = NULL, *p, keyFound = 0;
  unsigned ret = 4;
  kf = fopen(fname, "rb");
  if (kf == NULL) {
    perror("Keyfile open");
    return 1;
  }
  line = calloc(2341, 1); /* Maybe enough */
  while (fgets(line, 2341, kf)) {
    p = line;
    while (isspace(*p)) {
      p++;
    }
    /*
      Key in VSDSP format.
      |  key[0]     | key [1] |
      | LSW   | MSW |
      | M  L  |
      HEXKEY:  28 8e  cc68  b4ef 60ee 49a2 de1d 8f16 406b 609d c4e6 5494 6af2 0e9e 4e7c b952 94bd
      HEXKEY:288ecc68b4ef60ee49a2de1d8f16406b609dc4e654946af20e9e4e7cb95294bd
    */
    if (strncmp(p, "HEXKEY:", 7) == 0) {
      char *e;
      if (keyFound) {
	ret = 3;
	goto finally;
      }
      p += 7;
      e = p + strlen(p) - 1;
      while (isspace(*e)) {
	*e-- = 0;
      }
		
      if (strlen(p) == 64) {
	unsigned char keyBuf[32], *k, i;
	/* a good line */
	k = keyBuf;
	for (i = 0; i < 32; i++) {
	  if (isxdigit(*p)) {
	    if (isdigit(*p)) {
	      *k = ((*p - '0') << 4);
	    } else {
	      *k = ((toupper(*p) - 'A' + 10) << 4) ;
	    }
	    p++;
	  } else {
	    ret = 2;
	    goto finally;
	  }
	  if (isxdigit(*p)) {
	    if (isdigit(*p)) {
	      *k |= *p - '0';
	    } else {
	      *k |= toupper(*p) - 'A' + 10;
	    }
	    p++;
	  } else {
	    ret = 2;
	    goto finally;
	  }
	  k++;
	}
	k = keyBuf;
	for (i = 0; i < 8; i++) {
	  *key = ((((u_int32)(*k++)) & 0xff) << 8);
	  *key |= ((((u_int32)(*k++)) & 0xff) << 0);
	  *key |= ((((u_int32)(*k++)) & 0xff) << 24);
	  *key++ |= ((((u_int32)(*k++)) & 0xff) << 16);
	}
	keyFound = 1;
	ret = 0;
      }
    }
  }
 finally:
  if (kf) {
    fclose(kf);
  }
  if (line) {
    free(line);
  }
  return ret;
}


static int srandDone = 0;


/* Generate 256-bit key */
int GenerateKey(const char *fileName, const char *salt) {
  int saltEnthropy = 0, totEnthropy=0;
  u_int32 key[8] = {0};
  int i;
  int bitIdx = 0, wordIdx = 0, c;
  int retCode = -1;
  FILE *fp = NULL;
  
#if 1
  /*
    Start up nonce with current time of day.
    Actual enthropy of this is
    - In one day: log(3600*24)/log(2) = 16.4 bits
    - In one year: log(3600*24)/log(2) = 24.9 bits
    So, by itself, it is a very poor source for 256-bit encryption.
    However, it is a starting point.
   */
  if (!srandDone) {
    srandDone = 1;
    srand(time(NULL));
  }
  key[7] ^= time(NULL);
  if (RAND_MAX < 0x8000) {
    totEnthropy += 15;
  } else {
    totEnthropy += 25;
  }
#if 1
  for (i=0; i<8; i++) {
    if (RAND_MAX < 0x10000) {
      key[i] ^= rand() ^ (rand()<<16);
    } else {
      key[i] ^= rand();
    }
  }
#endif
#endif

#if 1
  /*
    Continue with adding extra enthropy.
    If tim.tv_nsec actually counts real nanoseconds, the enthropy
    added by this part is log(1e9)/log(2) = 29.9 bits for the
    nanosecond part, and none for the full seconds part when
    combined with the previous time(NULL) initialization. However,
    at minimum, if tv_nsec doesn't contain useful information, this
    part adds nothing to enthropy.
   */
  struct timespec tim = {0};
  if (!clock_gettime(CLOCK_REALTIME, &tim)) {
    key[6] ^= tim.tv_sec;
    key[5] ^= tim.tv_nsec;
  }
#endif
  
  /* The main enthropy comes from the salt. */
  if (salt) {
    saltEnthropy = EstimateEntropy((unsigned char *)salt, 0);
    totEnthropy += saltEnthropy;

#if 1
    while ((c=*salt++)) {
      key[wordIdx] ^= c<<bitIdx;
      if (bitIdx > 32-8) {
	key[(wordIdx+1)%8] ^= c>>(32-bitIdx);
      }

      bitIdx += 7;
      if (bitIdx >= 32) {
	bitIdx -= 32;
	if (++wordIdx >= 8) {
	  wordIdx = 0;
	}
      }
    }
#endif
  }
  
  printf("Estimated potential key enthropy %d bits: ", totEnthropy);
  if (totEnthropy >= 256) {
    printf("Good.\n");
  } else if (totEnthropy >= 128) {
    printf("Passable but could be better.\n");
  } else {
    printf("BAD! REFUSING TO WRITE!\n");
    goto finally;
  }

  fp = fopen(fileName, "r");
  if (fp) {
    fprintf(stderr, "File \"%s\" exists, refusing to overwrite!\n"
	    "  If you want a file by this name, first remove it by hand!\n",
	    fileName);
    goto finally;
  }
  fp = fopen(fileName, "w");
  if (!fp) {
    fprintf(stderr, "Couldn't open \"%s\" for writing\n", fileName);
    goto finally;
  }
  fprintf(fp, "HEXKEY:");
  printf("HEXKEY:");
  for (i=7; i>=0; i--) {
    fprintf(fp, "%08x", key[i]);
    printf("%08x", key[i]);
  }
  fprintf(fp, "\n");
  printf("\nKey written to file \"%s\".\n", fileName);
  
  retCode = 0;
 finally:
  if (fp) {
    fclose(fp);
    fp = NULL;
  }
  return retCode;
}

/* Generate 96-bit nonce from the file to be encoded.
   The better enthropy the file has, the better the generated nonce
   will be. */
int GenerateNonce(FILE *fp, u_int32 *nonce) {
  int retCode = -1;
  memset(nonce, 0, 3*sizeof(*nonce));
  int bitIdx = 0, wordIdx = 0, c;

#if 1
  /*
    Start up nonce with current time of day.
    Actual enthropy of this is
    - In one day: log(3600*24)/log(2) = 16.4 bits
    - In one year: log(3600*24)/log(2) = 24.9 bits
    So, by itself, it is a very poor source for 256-bit encryption.
    However, it is a starting point.
   */
  if (!srandDone) {
    srandDone = 1;
    srand(time(NULL));
  }
  if (RAND_MAX < 0x10000) {
    nonce[2] ^= time(NULL) ^ rand() ^ (rand()<<16);
    nonce[1] ^= rand() ^ (rand()<<16);
    nonce[0] ^= rand() ^ (rand()<<16);
  } else {
    nonce[2] ^= time(NULL) ^ rand();
    nonce[1] ^= rand();
    nonce[0] ^= rand();
  }
#endif

#if 1
  /*
    Continue with adding extra enthropy.
    If tim.tv_nsec actually counts real nanoseconds, the enthropy
    added by this part is log(1e9)/log(2) = 29.9 bits for the
    nanosecond part, and none for the full seconds part when
    combined with the previous time(NULL) initialization. However,
    at minimum, if tv_nsec doesn't contain useful information, this
    part adds nothing to enthropy.
   */
  struct timespec tim = {0};
  if (!clock_gettime(CLOCK_REALTIME, &tim)) {
    nonce[1] ^= tim.tv_sec;
    nonce[0] ^= tim.tv_nsec;
  }
#endif
  
#if 1
  /* The main enthropy comes from the file itself. If the file is really
     short, this does not help much and we are left with what we got from
     the previous stages. However, for e.g. audio files, this is pretty
     much certain to mingle all bits to very pseudo-random positions. */
  while ((c=fgetc(fp)) >= 0) {
    nonce[wordIdx] ^= c<<bitIdx;
    if (bitIdx > 32-8) {
      nonce[(wordIdx+1)%3] ^= c>>(32-bitIdx);
    }

    bitIdx += 7;
    if (bitIdx >= 32) {
      bitIdx -= 32;
      if (++wordIdx >= 3) {
	wordIdx = 0;
      }
    }
  }
  fseek(fp, 0, SEEK_SET);
#endif
  
#if 0
  printf("nonce: %08x %08x %08x\n", nonce[2], nonce[1], nonce[0]);
#endif

  retCode = 0;

  return retCode;
}


#if 0
void GenerateRandom(u_int32 *out, char *salt) {
  /* Not cryptographically good.
     Arc4Random() would be much better.
  */
  u_int32 key[8], counter;
  unsigned char *cp = (unsigned char *)key;
  static u_int32 inited = 0;
  int saltEnthropy;
  int useWords = out ? 3 : 8;
  
  if (salt) {
    saltEnthropy = EstimateEntropy((unsigned char *)salt, 0);
  }
  
  if (inited == 0) {
    /* time(NULL) enthropy:
       - In one day: log(3600*24)/log(2) = 16.4 bits
       - In one year: log(3600*24)/log(2) = 24.9 bits
       So, by itself, it is a very poor source for 256-bit encryption.
       However, it is a starting point. */
    srand(time(NULL));
    inited = 12;
  }

#if 0
  for (counter = 0; counter < 3; counter++) {
    key[counter] = rand();
    nonce[counter] = rand();
  }
  for (counter = 3; counter < 8; counter++) {
    key[counter] = rand();
  }

  Chacha20SetKey(key, counter, nonce);
  Chacha20Encrypt((char *)key, 32);
  if (out == NULL) {
    printf("HEXKEY:");
    for (counter = 0; counter < 32; counter++) {
      printf("%02hx", *cp++);
    }
    printf("\n");
  } else {
    out[0] = key[0];
    out[1] = key[1];
    out[2] = key[2];
  }
#endif
}
#endif


#define LINE_LEN 8192
int EncryptFiles(char *listName, u_int32 *key) {
  FILE *list=NULL, *fi=NULL, *fo=NULL;
  char *ofileName = NULL;
  char *ifileName = NULL;
  char *line = NULL;
  char readBuf[64];
  u_int32 nonce[3];
  int ret = -1;
  size_t n = 0;
  list = fopen(listName, "rb");
  if (list == NULL) {
    perror("listfile open");
    return 1;
  }
  line = calloc(LINE_LEN, 1);
  ifileName = calloc(LINE_LEN, 1);
  ofileName = calloc(LINE_LEN, 1);
  if (ifileName == NULL) {
    perror("ifilename allocation");
    ret = 5;
    goto finally;
  }
  while (fgets(line, LINE_LEN - 1, list)) {
    int inQuotes = 0;
    int i=0, o=0;
    n = strlen(line);
#if 0
    while (n && isspace(line[n - 1])) {
      line[--n] = 0;
    }
    ifilename = line;
    while (isspace(ifilename)) ifilename++;
    t = ifilename;
    while (!isspace(*t) || inQuotes) {
      if (!t) {
	fprintf(stderr, "ERROR: Malformed line \"%s\"\n", ifilename);
      }
      
    }
#endif
    while (isspace(line[i])) i++;
    while (line[i] && (inQuotes || !isspace(line[i]))) {
      if (line[i] == '"') {
	inQuotes = !inQuotes;
      } else {
	ifileName[o++] = line[i];
      }
      i++;
    }
    ifileName[o++] = '\0';
    o = 0;
    while (isspace(line[i])) i++;
    while (line[i] && (inQuotes || !isspace(line[i]))) {
      if (line[i] == '"') {
	inQuotes = !inQuotes;
      } else {
	ofileName[o++] = line[i];
      }
      i++;
    }
    while (isspace(line[i])) i++;
    ofileName[o++] = '\0';
    if (line[i] || inQuotes || strlen(ifileName) == 0 ||
	strlen(ofileName) == 0) {
      fprintf(stderr, "Malformed line in %s:\n%s\n", listName, line);
      ret = 7;
      goto finally;
    }
    
    fi = fopen(ifileName, "rb");
    if (fi == NULL) {
      fprintf(stderr, "ERROR: Input file: '%s'\n", ifileName);
      perror(ifileName);
      ret = 2;
      goto finally;
    }
    printf("Encrypt '%s' to '%s'\n", ifileName, ofileName);
    
    fo = fopen(ofileName, "wb");
    if (fo == NULL) {
      perror(ofileName);
      ret = 2;
      goto finally;
    }
    
    GenerateNonce(fi, nonce);
    Chacha20SetKey(key, 0, nonce);
    while (feof(fi) == 0) {
      n = fread(readBuf, 1, 64, fi);
      if (ferror(fi)) {
	perror(ifileName);
	ret = 5;
	goto finally;
      }
      Chacha20Encrypt(readBuf, n);
      if (fwrite(readBuf, 1, n, fo) != n) {
	perror(ofileName);
	ret = 6;
	goto finally;
      }
    }
    fclose(fi);
    fi = NULL;
		

    for (n = 0; n < 3; n++) {
      readBuf[(n * 4) + 0] = (nonce[n] >> 8) & 0xff;
      readBuf[(n * 4) + 1] = (nonce[n] >> 0) & 0xff;
      readBuf[(n * 4) + 2] = (nonce[n] >> 24) & 0xff;
      readBuf[(n * 4) + 3] = (nonce[n] >> 16) & 0xff;
    }
#if 0
    printf("nonce: ");
    for (n= 0; n < 12; n++) {
      printf("%02hx ", readBuf[n] & 0xff);
    }
    printf("\n%08x %08x %08x\n", nonce[0], nonce[1], nonce[2]);
#endif		
    if (fwrite(readBuf, 1, 12, fo) != 12) {
      perror("Crypt write nonce");
      ret = 4;
      goto finally;
    }
    fclose(fo);
    fo = NULL;
  }
  ret = 0;
 finally:
  if (ifileName) {
    free(ifileName);
  }
  if (ofileName) {
    free(ofileName);
  }
  if (line) {
    free(line);
  }
  if (list) {
    fclose(list);
  }
  if (fi) {
    fclose(fi);
  }
  if (fo) {
    fclose(fo);
  }
  return ret;
}

int main(int argc , char *argv[]) {
  char *generateKeyTemplate = NULL;
  int i;
  u_int32 trueKey[8];
  char *keyFile = NULL, *inputFile = NULL, **nameTarget = NULL;
  int retCode = EXIT_FAILURE;
  char *salt = NULL;

  for (i = 1; i < argc; i++) {
    if (nameTarget) {
      *nameTarget = argv[i];
      nameTarget = NULL;
    } else if (strcmp(argv[i], "-h") == 0) {
      printf("VLSIEncryptor 2023-11-14\n"
	     "Usage:\n"
	     "%s [-h] [-s salt -g keyfile | -k keyfile -i inlistfile]\n"
	     "-s salt\tProvide random salt data\n"
	     "\tTo create good salt data, push random keys to make salt"
	     " enthropy\n"
	     "\tgood enough. Example:\n"
	     "\thrfyuntnkwtnnyt5t4n9t83tm3o2mtcy3m4tm34tcmttmcu8t49y8mv893t3J"
	     "IFEOW\n"
	     "\tDO NOT COPY THIS AS IS BUT MASH YOUR OWN ON YOUR KEYBOARD!\n"
	     "-g keyf\tGenerate secret key file keyf (e.g. my.key)\n"
	     "-k keyf\tUse secret key file keyf for encryption (e.g. my.key)\n"
	     "-i ilf\tEncrypt files from inlistfile\n"
	     "\tEach line in inlistfile should look like the following:\n"
	     "\thello.mp3 hellomp3.crp\n"
	     "\tThe previous line tells to encrypt hello.mp3 to hellomp3.crp\n"
	     "\tBe careful not to overwrite source files: the file names on\n"
	     "\tthe line must be different!\n"
	     "-h\tShow this help\n"
	     , argv[0]);
      return 0;
    } else if (strcmp(argv[i], "-s") == 0) {
      if (i >= argc-1) {
	fprintf(stderr, "%s: salt data for parameter '-s' missing\n",
		argv[0]);
	goto finally;
      }
      salt = argv[++i];
    } else if (strcmp(argv[i], "-g") == 0) {
      if (i >= argc-1) {
	fprintf(stderr, "%s: file name for parameter '-g' missing\n",
		argv[0]);
	goto finally;
      }
      generateKeyTemplate = argv[++i];
    } else if (strcmp(argv[i], "-i") == 0) {
      nameTarget = &inputFile;
    } else if (strcmp(argv[i], "-k") == 0) {
      nameTarget = &keyFile;
    } else {
      fprintf(stderr, "%s: unkonwn parameter '%s'\n", argv[0], argv[i]);
      return -1;
    }
  }

  if (generateKeyTemplate) {
    if (GenerateKey(generateKeyTemplate, salt)) {
      fprintf(stderr, "%s: Failed to create key file \"%s\"\n",
	      argv[0], generateKeyTemplate);
      goto finally;
    }
    retCode = EXIT_SUCCESS;
    goto finally;
  }
  if (keyFile) {
    unsigned r;
    const char *values[] = {
			    "OK",
			    "",/* Open, perror()ed */
			    "Value problem",
			    "Multiple keys",
			    "No key"
    };
    printf("Reading key");
    r = ReadKey(keyFile, trueKey);
    if (r < (sizeof(values) / sizeof(values[0]))) {
      printf(" %s\n", values[r]);
      if (r) {
	fprintf(stderr, "KEY ERROR\n");
	return -1;
      }
#if 0
      else {
	for (r = 0; r < 8; r++) {
	  printf("%08x ", trueKey[r]);
	}
	printf("\n");
      }
#endif
    } else {
      fprintf(stderr, "%s: program error read key %u\n", argv[0], r);
      return -1;
    }
  } else {
    fprintf(stderr, "%s: Key file not defined\n", argv[0]);
    return -1;
  }
  if (inputFile) {
    return  EncryptFiles(inputFile, trueKey);
  } else {
    fprintf(stderr, "%s: Input file not defined\n", argv[0]);
    return -1;
  }

  retCode = EXIT_SUCCESS;
 finally:
  return retCode;
}
