/* %W% (Qualcomm) %G% */
/*
 * Encrypt an ASCII file, leaving it in ASCII, without expansion.
 * Uses Arrsyfor (compatible with RC4) as its base cipher.
 */

#include <stdio.h>
#include <time.h>
#include <string.h>

/*
 * arrsyfor included:
 * Code written by Greg Rose, based on the textual description in
 * Schneier's AC2, pp 397-398. Intended to be compatible with RC4,
 * except that the first 256 bytes of output are discarded after key
 * setup.
 *
 * No rights reserved.
 */

typedef struct {
	unsigned char	Sbox[256];
	unsigned char	i, j;
} ctx_arrsyfor;

ctx_arrsyfor	c;

#define S (c.Sbox)
#define SWAP(i,j)   {\
			unsigned char	t;\
			t = S[i];\
			S[i] = S[j];\
			S[j] = t;\
		    }

/*
 * Encrypt buffer by exclusive-or with generated octets. If you really want
 * to see the actual bytes, zero a buffer first.
 */
void
arrsyfor(unsigned char *buf, register int n)
{
    register int	i, j;

    i = c.i; j = c.j;
    while (--n >= 0) {
	i = (i + 1) & 0xFF;
	j = (j + S[i]) & 0xFF;
	SWAP(i, j);
	*buf++ ^= S[(S[i]+S[j]) & 0xFF];
    }
    c.i = i; c.j = j;
}

/*
 * initialise the arrsyfor context, including key setup and initial
 * output disposal (Kocher, I believe).
 */
void
key_arrsyfor(unsigned char *key, int n)
{
    register int	i;
    unsigned char	j;
    unsigned char	tbuf[256];

    for (i = 0; i < 256; ++i)
	S[i] = i;
    j = 0;
    for (i = 0; i < 256; ++i) {
	j += S[i] + key[i % n];
	SWAP(i, j);
    }
    c.i = c.j = 0;
    /* burn the first 256 outputs, they aren't as random as one would like */
    arrsyfor(tbuf, 256);
}

unsigned char
onebyte()
{
    unsigned char	t = 0;
    arrsyfor(&t, 1);
    return t;
} 

/*
 * The following string contains all the characters that will be mapped
 * into other characters in the set. Characters not in the string will be
 * left alone. The function "setupalphabet" creates an inverse table from
 * this string.
 */
char	alphabet[] =
	    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	    "abcdefghijklmnopqrstuvwxyz"
	    "0123456789";
#define N (sizeof(alphabet) - 1)
#define MAXN (N * ((255 + N)/N))
unsigned char	revalpha[256];

void
setupalphabet()
{
    register int	i;

    for (i = 0; i < 256; ++i)
	revalpha[i] = 0xFF;
    for (i = 0; i < N; ++i)
	revalpha[(unsigned char)alphabet[i]] = i;
}

/*
 * encrypt a single character and output it.
 * Can never be made to encrypt NUL
 * but can actually handle other high-bit-set characters.
 * We want it to be self-inverse, so it's a bit more complicated
 * than straight RC4. Basically, we add a random offset, pass through
 * an involution operation, and subtract the offset again. So long as the
 * random offsets line up, we'll be able to decrypt the same way.
 * Since the involution is not cryptographically significant, we just
 * reflect the accepted alphabet.
 */
void
encrypt(int c)
{
    register int	b, r;

    if ((r = revalpha[c] & 0xFF) != 0xFF) {
	do {
	    b = onebyte();
	} while (b >= MAXN);
	b %= N;
	c = ((2*N - 1) - ((r + b) % N) - b) % N;
	putchar(alphabet[c]);
    }
    else
	putchar(c);
}

#define MAX	100	/* maximum length of various stuff */

char	*myname = "asciirc4";
char	*key, *IV;
int	keylen;
char	buffer[MAX];
char	IVbuf[MAX];
char	keybuf[MAX];
char	autostr[] = "asciirc4 IV: ";
int	autolen;
int	autofound = 0;

/*
 * arguments: key [IV]
 */
int
main(int ac, char **av)
{
    register int	i;

    if (av[0] != NULL)
	myname = av[0];
    ++av, --ac;

    /* figure out arguments */
    if (ac != 1 && ac != 2) {
      usage:
	fprintf(stderr, "usage: %s ascii-key [hex-IV]\n", myname);
	fprintf(stderr, "\twill try to auto-detect IV if not given\n");
	return 1;
    }

    if (ac >= 1) {
	key = av[0];
	keylen = strlen(key);
    }

    /* try to autodetect an IV */
    if (fgets(buffer, MAX, stdin) == NULL) {
	fprintf(stderr, "Error attempting to read IV\n");
	return 2;
    }
    else {
	/* if it's the correct format for an IV string, use it. Otherwise
	 * leave it lying around to become the first input.
	 */
	if (strncmp(buffer, autostr, strlen(autostr)) == 0
	     && buffer[strlen(buffer)-1] == '\n') {
	    buffer[strlen(buffer) - 1] = '\0';
	    IV = buffer + strlen(autostr);
	    autofound = 1;
	    autolen = 0;
	}
	else
	    autolen = strlen(buffer);
    }

    if (ac == 2) {
	/* what if also found one in the encryption? It's hard
	 * to know what's right. If they're equal, silently drop it.
	 * Otherwise, complain... something funny is going on.
	 */
	if (autofound && strcmp(IV, av[1]) != 0) {
	    fprintf(stderr, "Different IVs found and specified.\n");
	    return 5;
	}
	IV = av[1];
    }
    else if (!autofound) {
	/* make up an IV and output it in appropriate form */
	sprintf(IVbuf, "%08lx", time(NULL));
	IV = IVbuf;
	printf("%s%s\n", autostr, IV);
    }

    /* At this point we have a key and an initialisation vector.
     * Rivest recommends hashing them together and using the output, but
     * I don't have a hash function lying around. Arrsyfor itself can be
     * used though...
     * 1. Key it with the concatenation of key and IV
     * 2. Discard 256 bytes
     * 3. Take 32 bytes for new key
     * 4. Discard 256 bytes
     * ... then use it
     */
    keybuf[0] = '\0';
    strncat(keybuf, key, MAX);
    strncat(keybuf, IV, MAX - strlen(key) - 1);
    if (strlen(keybuf) > MAX-1) {
	fprintf(stderr, "Combination of key and IV is too long (max %d)\n",
	    MAX-1);
	goto usage;
    }
    key_arrsyfor((unsigned char *)keybuf, strlen(keybuf));
    for (i = 0; i < MAX; ++i)
	keybuf[i] = 0;
    arrsyfor((unsigned char *)keybuf, 32);
    key_arrsyfor((unsigned char *)keybuf, 32);

    setupalphabet();

    /* now perform the encryption, first using up any bytes read
     * while looking for an IV.
     */
    for (i = 0; i < autolen; ++i)
	encrypt(buffer[i]);
    while ((i = getchar()) != EOF)
	encrypt(i);
    return 0;
}