//CryptStream.h

#if _MSC_VER > 1000
#pragma once
#endif // _MSC_VER > 1000

#ifndef CryptStreamh
#define CryptStreamh

#include "RC4.h"

/*
XOR256 Stream Encryption:
Encrypt [in chunks].
Reset.
Decrypt [in chunks].
  CCryptStream Stream;
  if(!Stream.SetDepth(12)) return;
  if(!Stream.SetKey("Key")) return;
  char Text[]="This is the text";
  int Len=strlen(Text);
  int Len1=Len>>1;
  int Len2=Len-Len1;

  Stream.Encrypt((BYTE*)Text,Text,Len1);
  Stream.Encrypt((BYTE*)Text+Len1,Text+Len1,Len2);

  Stream.Reset();
  Stream.Decrypt(Text,(BYTE*)Text,Len);
*/
class CCryptStream {
protected:
  CRC4 RC4;
  int  Depth;
  int  Depth1; //Depth-1
  bool Initialised;
  bool HasKey;
  bool HasDepth;
  BYTE* SeqX0;
  BYTE* SeqM0;
  BYTE Prev0;
  BYTE Prev;

public:
  CCryptStream() : SeqX0(0), SeqM0(0), Initialised(false), HasKey(false), HasDepth(false) {}
  CCryptStream(const char* Key, int _Depth) {Set(Key,_Depth);}

  bool Set(const char* Key, int _Depth) {return SetKey(Key) && SetDepth(_Depth);}

  bool SetDepth(int _Depth) {
    if(_Depth<1) return false;
    Depth=_Depth;
    Depth1=Depth-1;
    SeqX0=new BYTE[Depth];
    SeqM0=new BYTE[Depth];
    HasDepth=true;
    if(HasKey) Initialised=true;
    return true;
  }

  bool SetKey(const char* Key) {
    if(!*Key) return false;
    RC4.Seed(Key);
    Prev=Prev0=RC4.Rand();
    HasKey=true;
    if(HasDepth) Initialised=true;
    return true;
  }

  virtual ~CCryptStream() {
    delete[] SeqX0;
    delete[] SeqM0;
  }

  bool Reset() {
    bool Ret=RC4.Reset();
    RC4.Rand();
    Prev=Prev0;
    return Ret;
  }

  bool Encrypt(BYTE* dst, const char* src, unsigned Len) {
    if(!Initialised) return false;
    for(unsigned i=0; i<Len; ++i) {
      //The first round
      Prev ^= *src++;
      BYTE b=(Prev ^ RC4.Rand()) + RC4.Rand();
      //The last Depth-1 Depth
      for(int j=1; j<Depth; ++j) {
        b ^= RC4.Rand();
        b += RC4.Rand();
      }
      *dst++=b;
    }
    return true;
  }

  bool Decrypt(char* dst, const BYTE* src, unsigned Len) {
    if(!Initialised) return false;
    for(unsigned i=0; i<Len; ++i) {
      //Calculate the constants
      BYTE* SeqX=SeqX0;
      BYTE* SeqM=SeqM0;
      for(int j=0; j<Depth; ++j) {
        *(SeqX++)=RC4.Rand();
        *(SeqM++)=RC4.Rand();
      }
      //The last Depth-1 Depth
      BYTE Plain=(BYTE)*src++;
      SeqX=SeqX0 + Depth1;
      SeqM=SeqM0 + Depth1;
      for(j=Depth-1; j>0; --j,--SeqX,--SeqM) {
        if(*SeqM <= Plain) Plain -= *SeqM;
        else (Plain += ~*SeqM)++;
        Plain ^= *SeqX;
      }
      //First round
      if(*SeqM0 <= Plain) Plain -= *SeqM0;
      else (Plain += ~*SeqM0)++;
      Plain ^= Prev ^ *SeqX0;
      *dst++=Plain;
      Prev ^= Plain;
    }
    return true;
  }
};

#endif //CryptStreamh