#include <string.h>
#include <string>

using namespace std;

const unsigned char Se[16] = {	0xE, 0x4, 0xD, 0x1,
				0x2, 0xF, 0xB, 0x8,
				0x3, 0xA, 0x6, 0xC,
				0x5, 0x9, 0x0, 0x7};

const unsigned char Sd[16] = {	0xE, 0x3, 0x4, 0x8,
				0x1, 0xC, 0xA, 0xF,
				0x7, 0xD, 0x9, 0x6,
				0xB, 0x2, 0x0, 0x5};

const unsigned char Rounds = 4;

typedef unsigned char TKey[8];
typedef unsigned char TSubKey[2];
typedef unsigned char TMessageBlock[2];
typedef TMessageBlock* PMessageBlock;
typedef unsigned char* ptr;

class CWeakHaze16 {
  private:
    TSubKey FSubKeys[5];
    TKey FKey;

    void CreateSubKeys();
    void SubKeyMixing(PMessageBlock m, unsigned char RoundNum);
    void SBoxes(PMessageBlock m, bool encrypting);
    void Permutation(PMessageBlock m);

    void EncRound(PMessageBlock m, unsigned char RoundNum);
    void DecRound(PMessageBlock m, unsigned char RoundNum);

  public:
    CWeakHaze16();
    
    void SetKey(TKey &newKey);
    void Encrypt(ptr inData, ptr outData, unsigned long int size);
    void Decrypt(ptr inData, ptr outData, unsigned long int size);
    
    short int SelfTest();
};

CWeakHaze16::CWeakHaze16() { }

void CWeakHaze16::CreateSubKeys() {
  unsigned short int i; 

  for (i=0;i<8;i++) 
    FSubKeys[i / 2][i % 2] = FKey[i];

  for (i=0;i<2;i++)
    FSubKeys[4][i] = FSubKeys[0][1 - i] ^ FSubKeys[1][i] ^
      FSubKeys[2][i] ^ FSubKeys[3][1 - i];
}

void CWeakHaze16::SubKeyMixing(PMessageBlock m, unsigned char RoundNum) {
  for (unsigned short int i = 0; i < 2; i++)
    (*m)[i] = (*m)[i] ^ FSubKeys[RoundNum-1][i];
}

void CWeakHaze16::SBoxes(PMessageBlock m, bool encrypting) {
  if (encrypting) {
    for (unsigned short int i = 0; i < 2; i++)
      (*m)[i] = (Se[(*m)[i] >> 4] << 4) ^ Se[(*m)[i] & 0xf];
  } else {
    for (unsigned short int i = 0; i < 2; i++)
      (*m)[i] = (Sd[(*m)[i] >> 4] << 4) ^ Sd[(*m)[i] & 0xf];
  }
}

void CWeakHaze16::Permutation(PMessageBlock m) {
  unsigned char t = (*m)[0];
  (*m)[0] = ((*m)[0] & 0x80) ^ (((*m)[0] << 3) & 0x40) ^ // bits 1 + 2
    (((*m)[1] >> 2) & 0x20) ^ (((*m)[1] << 1) & 0x10) ^ // bits 3 + 4
    (((*m)[0] >> 3) & 0x8) ^ ((*m)[0] & 0x4) ^ // bits 5 + 6
    (((*m)[1] >> 5) & 0x2) ^ (((*m)[1] >> 2) & 0x1); // bits 7 + 8

  (*m)[1] = ((t << 2) & 0x80) ^ ((t << 5) & 0x40) ^ // bits 1 + 2
    ((*m)[1] & 0x20) ^ (((*m)[1] << 3) & 0x10) ^ // bits 3 + 4
    ((t >> 1) & 0x8) ^ ((t << 2) & 0x4) ^ // bits 5 + 6
    (((*m)[1] >> 3) & 0x2) ^ ((*m)[1] & 0x1); // bits 7 + 8
}

void CWeakHaze16::EncRound(PMessageBlock m, unsigned char RoundNum) {
  if (RoundNum < Rounds) {
    SubKeyMixing(m, RoundNum);
    SBoxes(m, true);
    Permutation(m);
  } else {
    SubKeyMixing(m, RoundNum);
    SBoxes(m, true);   
  }
}

void CWeakHaze16::DecRound(PMessageBlock m, unsigned char RoundNum) {
  if (RoundNum < Rounds) {
    Permutation(m);
    SBoxes(m, false);
    SubKeyMixing(m, RoundNum);
  } else {
    SBoxes(m, false);
    SubKeyMixing(m, RoundNum); 
  }
}

void CWeakHaze16::SetKey(TKey &newKey) {
  memcpy(FKey, newKey, 8); 
//  FKey = newKey;
  CreateSubKeys();
}

void CWeakHaze16::Encrypt(ptr inData, ptr outData, unsigned long int size) {
  PMessageBlock m;
  TMessageBlock tm;
  unsigned long int i;
  unsigned short int j;

  memcpy(outData, inData, size);

  m = (PMessageBlock)outData;

  i = 0;
  while (i++ < size / 2) {
    for (short int j = 1; j < 5; j++)
      EncRound(m, j);
    SubKeyMixing(m, 5);
    m++; 
  }

  if (size % 2 == 1) {
    tm[0] = 0;
    tm[1] = 0; // yes, this should lead to a known plaintext attack.
    for (j = 1; j < 5; j++)
      EncRound(&tm, j);
    SubKeyMixing(&tm, 5);

    (*m)[0] = tm[0] ^ (*m)[0];
  }
}

void CWeakHaze16::Decrypt(ptr inData, ptr outData, unsigned long int size) {
  PMessageBlock m;
  TMessageBlock tm;
  unsigned long int i;
  unsigned short int j;

  memcpy(outData, inData, size);

  m = (PMessageBlock)outData;

  i = 0;
  while (i++ < size / 2) {
    SubKeyMixing(m, 5);
    for (short int j = 4; j > 0; j--)
      DecRound(m, j);
    m++; 
  }

  if (size % 2 == 1) {
    tm[0] = 0;
    tm[1] = 0; // yes, this should lead to a known plaintext attack.
    for (j = 1; j < 5; j++)
      EncRound(&tm, j);
    SubKeyMixing(&tm, 5);

    (*m)[0] = tm[0] ^ (*m)[0];
  }
}

short int CWeakHaze16::SelfTest() { // returns 0 on success, else test that failed.
  TKey k;
  unsigned short int i;
  TMessageBlock m;

  for(int i=0; i<8;i++)
    k[i] = 0xff;

  SetKey(k);
  m[0] = 0;
  m[1] = 0;
  
  Encrypt(m, m, 2);
  if (m[0] != 0x3a)
    return 1;
  if (m[1] != 0x48)
    return 2;

  Decrypt(m, m, 2);
  if (m[0] != 0)
    return 3;
  if (m[1] != 0)
    return 4;

  char s[] = "testing!";
  Encrypt(s, s, strlen(s));
  Decrypt(s, s, strlen(s));
  if (strcmp(s, "testing!") != 0)
    return 5;

  char s2[] = "testing";
  Encrypt(&s, &s, strlen(s));
  Decrypt(&s, &s, strlen(s));
  if (strcmp(s, "testing") != 0)
    return 6;

  // haha. these tests suck, they should have literals. :)
  SubKeyMixing(&m, 1);
  SubKeyMixing(&m, 2);
  SubKeyMixing(&m, 3);
  SubKeyMixing(&m, 4);
  SubKeyMixing(&m, 5);
  SubKeyMixing(&m, 1);
  SubKeyMixing(&m, 2);
  SubKeyMixing(&m, 3);
  SubKeyMixing(&m, 4);
  SubKeyMixing(&m, 5);
  if ((m[0] != 0) || (m[1] != 0)) 
    return 7;

  return 0; // if you made it here, you passed.

}