#include "LIB_NG_NAS_Functions.hh"

#include "loggers.hh"

#include "rijndael.hh"
#include "opc.hh"

namespace LIB__NG__NAS__Functions {

  static uint8_t OP[16] = {0}; // FIXME FSCOM To be refined. This is a Q&D implementation

  void fx__set__op(const OCTETSTRING& p_op) {
    loggers::get_instance().log_msg(">>> fx__set__op: p_op: ", p_op);
    std::memcpy(OP, static_cast<const unsigned char*>(p_op), 16);
    loggers::get_instance().log_to_hexa("<<< fx__set__op: OP: ", static_cast<const unsigned char*>(OP), 16);
  }

  INTEGER fx__f1(const BITSTRING& p_authK, const BITSTRING& p_rand, const BITSTRING& p_sqn, const BITSTRING& p_amf, BITSTRING& p_mac_a) {
    loggers::get_instance().log_msg(">>> fx__f1: p_authK: ", bit2oct(p_authK));
    loggers::get_instance().log_msg(">>> fx__f1: p_rand: ", bit2oct(p_rand));

    rijndael r;
    OCTETSTRING authK = bit2oct(p_authK);
    r.rijndael_key_schedule(authK);
    opc op(r, OP);
    uint8_t op_c[16] = { 0x00 };
    op.compute_opc(op_c);
  
    OCTETSTRING rand = bit2oct(p_rand);
    uint8_t rijndael_input[16] = { 0x00 };
    for (int i = 0; i < 16; i++) {
      rijndael_input[i] = rand[i].get_octet() ^ op_c[i];
    } // End of 'for' statement
    uint8_t temp[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, temp);

    OCTETSTRING sqn = bit2oct(p_sqn);
    uint8_t in1[16] = { 0x00 };
    for (int i = 0; i < 6; i++) {
      in1[i] = sqn[i].get_octet();
      in1[i + 8] = sqn[i].get_octet();
    } // End of 'for' statement
    OCTETSTRING amf = bit2oct(p_amf);
    for (int i = 0; i < 2; i++) {
      in1[i + 6] = amf[i].get_octet();
      in1[i + 14] = amf[i].get_octet();
    } // End of 'for' statement

    // XOR op_c and in1, rotate by r1=64, and XOR on the constant c1 (which is all zeroes)
    for (int i = 0; i < 16; i++) {
      rijndael_input[(i + 8) % 16] = in1[i] ^ op_c[i];
    }

    // XOR on the value temp computed before
    for (int i = 0; i < 16; i++) {
      rijndael_input[i] ^= temp[i];
    } // End of 'for' statement

    uint8_t out1[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, out1);
    for (int i = 0; i < 16; i++) {
      out1[i] ^= op_c[i];
    } // End of 'for' statement

    uint8_t mac_a[8] = { 0x00 };
    for (int i = 0; i < 8; i++) {
      mac_a[i] = out1[i];
    } // End of 'for' statement
    OCTETSTRING os(8, static_cast<const unsigned char*>(&mac_a[0]));
    p_mac_a = oct2bit(os);
    loggers::get_instance().log_msg("fx__f1star: p_mac_a: ", os);

    return 0;
  }

  INTEGER fx__f1star(const BITSTRING& p_authK, const BITSTRING& p_rand, const BITSTRING& p_sqn, const BITSTRING& p_amf, BITSTRING& p_mac_s) {
    loggers::get_instance().log_msg(">>> fx__f1star: p_authK: ", bit2oct(p_authK));
    loggers::get_instance().log_msg(">>> fx__f1star: p_rand: ", bit2oct(p_rand));

    rijndael r;
    OCTETSTRING authK = bit2oct(p_authK);
    r.rijndael_key_schedule(authK);
    opc op(r, OP);
    uint8_t op_c[16] = { 0x00 };
    op.compute_opc(op_c);
  
    OCTETSTRING rand = bit2oct(p_rand);
    uint8_t rijndael_input[16] = { 0x00 };
    for (int i = 0; i < 16; i++) {
      rijndael_input[i] = rand[i].get_octet() ^ op_c[i];
    } // End of 'for' statement
    uint8_t temp[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, temp);

    OCTETSTRING sqn = bit2oct(p_sqn);
    uint8_t in1[16] = { 0x00 };
    for (int i = 0; i < 6; i++) {
      in1[i] = sqn[i].get_octet();
      in1[i + 8] = sqn[i].get_octet();
    } // End of 'for' statement
    OCTETSTRING amf = bit2oct(p_amf);
    for (int i = 0; i < 2; i++) {
      in1[i + 6] = amf[i].get_octet();
      in1[i + 14] = amf[i].get_octet();
    } // End of 'for' statement

    // XOR op_c and in1, rotate by r1=64, and XOR on the constant c1 (which is all zeroes)
    for (int i = 0; i < 16; i++) {
      rijndael_input[(i + 8) % 16] = in1[i] ^ op_c[i];
    }

    // XOR on the value temp computed before
    for (int i = 0; i < 16; i++) {
      rijndael_input[i] ^= temp[i];
    } // End of 'for' statement

    uint8_t out1[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, out1);
    for (int i = 0; i < 16; i++) {
      out1[i] ^= op_c[i];
    } // End of 'for' statement

    uint8_t mac_s[8] = { 0x00 };
    for (int i = 0; i < 8; i++) {
      mac_s[i] = out1[i + 8];
    } // End of 'for' statement
    OCTETSTRING os(8, static_cast<const unsigned char*>(&mac_s[0]));
    p_mac_s = oct2bit(os);
    loggers::get_instance().log_msg("fx__f1star: p_mac_s: ", os);

    return 0;
  }

  INTEGER fx__f2345(const BITSTRING& p_authK, const BITSTRING& p_rand, BITSTRING& p_res, BITSTRING& p_ck, BITSTRING& p_ik, BITSTRING& p_ak) {
    loggers::get_instance().log_msg(">>> fx_f2345: p_authK: ", p_authK);
    loggers::get_instance().log_msg(">>> fx_f2345: p_rand: ", p_rand);

    rijndael r;
    OCTETSTRING authK = bit2oct(p_authK);
    r.rijndael_key_schedule(authK);
    opc op(r, OP);
    uint8_t op_c[16] = { 0x00 };
    op.compute_opc(op_c);
  
    OCTETSTRING rand = bit2oct(p_rand);
    uint8_t rijndael_input[16] = { 0x00 };
    for (int i = 0; i < 16; i++) {
      rijndael_input[i] = rand[i].get_octet() ^ op_c[i];
    } // End of 'for' statement
    uint8_t temp[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, temp);

    // To obtain output block OUT2: XOR OPc and TEMP, rotate by r2=0, and XOR on the constant c2 (which is all zeroes except that the last bit is 1)
    for (int i = 0; i < 16; i++) {
	    rijndael_input[i] = temp[i] ^ op_c[i];
    } // End of 'for' statement
    rijndael_input[15] ^= 1;

    uint8_t out[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, out);
    for (int i = 0; i < 16; i++) {
	    out[i] ^= op_c[i];
    } // End of 'for' statement

    uint8_t res[8] = { 0x00 };
    for (int i = 0; i < 8; i++) {
	    res[i] = out[i + 8];
    } // End of 'for' statement

    uint8_t ak[6] = { 0x00 };
    for (int i = 0; i < 6; i++) {
	    ak[i] = out[i];
    } // End of 'for' statement

    // To obtain output block OUT3: XOR OPc and TEMP, rotate by r3=32, and XOR on the constant c3 (which is all zeroes except that the next to last bit is 1)
    for (int i = 0; i < 16; i++) {
	    rijndael_input[(i + 12) % 16] = temp[i] ^ op_c[i];
    } // End of 'for' statement
    rijndael_input[15] ^= 2;

    r.rijndael_encrypt(rijndael_input, out);
    for (int i = 0; i < 16; i++) {
	    out[i] ^= op_c[i];
    } // End of 'for' statement

    uint8_t ck[16] = { 0x00 };
    for (int i = 0; i < 16; i++) {
	    ck[i] = out[i];
    } // End of 'for' statement

    // To obtain output block OUT4: XOR OPc and TEMP, rotate by r4=64, and XOR on the constant c4 (which is all zeroes except that the 2nd from last bit is 1)
    for (int i = 0; i < 16; i++) {
	    rijndael_input[(i+8) % 16] = temp[i] ^ op_c[i];
    } // End of 'for' statement
    rijndael_input[15] ^= 4;

    r.rijndael_encrypt(rijndael_input, out);
    for (int i = 0; i < 16; i++) {
      out[i] ^= op_c[i];
    } // End of 'for' statement

    uint8_t ik[16] = { 0x00 };
    for (int i = 0; i < 16; i++) {
      ik[i] = out[i];
    } // End of 'for' statement

    OCTETSTRING os(8, static_cast<const unsigned char*>(&res[0]));
    p_res = oct2bit(os);
    os = OCTETSTRING(16, static_cast<const unsigned char*>(&ik[0]));
    p_ik = oct2bit(os);
    os = OCTETSTRING(16, static_cast<const unsigned char*>(&ck[0]));
    p_ck = oct2bit(os);
    os = OCTETSTRING(6, static_cast<const unsigned char*>(&ak[0]));
    p_ak = oct2bit(os);
    loggers::get_instance().log_msg("fx_f2345: p_res: ", p_res);
    loggers::get_instance().log_msg("fx_f2345: p_ck: ", p_ck);
    loggers::get_instance().log_msg("fx_f2345: p_ik: ", p_ik);
    loggers::get_instance().log_msg("fx_f2345: p_ak: ", p_ak);

    return 0;
  }

  INTEGER fx__f5star(const BITSTRING& p_authK, const BITSTRING& p_rand, BITSTRING& p_ak) {
    loggers::get_instance().log_msg(">>> fx__f5star: p_authK: ", p_authK);
    loggers::get_instance().log_msg(">>> fx__f5star: p_rand: ", p_rand);

    rijndael r;
    OCTETSTRING authK = bit2oct(p_authK);
    r.rijndael_key_schedule(authK);
    opc op(r, OP);
    uint8_t op_c[16] = { 0x00 };
    op.compute_opc(op_c);
  
    OCTETSTRING rand = bit2oct(p_rand);
    uint8_t rijndael_input[16] = { 0x00 };
    for (int i = 0; i < 16; i++) {
      rijndael_input[i] = rand[i].get_octet() ^ op_c[i];
    } // End of 'for' statement
    uint8_t temp[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, temp);

    // To obtain output block OUT5: XOR OPc and TEMP, rotate by r5=96, and XOR on the constant c5 (which is all zeroes except that the 3rd from last bit is 1)
    for (int i = 0; i < 16; i++) {
      rijndael_input[(i + 4) % 16] = temp[i] ^ op_c[i];
    } // End of 'for' statement
    rijndael_input[15] ^= 8;

    uint8_t out[16] = { 0x00 };
    r.rijndael_encrypt(rijndael_input, out);
    for (int i = 0; i < 16; i++) {
      out[i] ^= op_c[i];
    } // End of 'for' statement

    uint8_t ak[6] = { 0x00 };
    for (int i = 0; i < 6; i++) {
      ak[i] = out[i];
    }
    OCTETSTRING os(6, static_cast<const unsigned char*>(&ak[0]));
    p_ak = oct2bit(os);
    loggers::get_instance().log_msg("fx__f5star: p_ak: ", os);

    return 0;
  }

} // End of namespace LIB__NG__NAS__Functions
