cryptx

// (c) Yura Zhivaga <yzhivaga@gmail.com>
// 2018

#include <iostream>
#include <vector>
#include <random>
#include <climits>
#include <algorithm>
#include <functional>
#include <boost/multiprecision/cpp_int.hpp>
#include <boost/random.hpp>

#include "sha/sha512.hh"

using namespace boost::random;
using namespace boost::multiprecision;
using Int = boost::multiprecision::number<boost::multiprecision::cpp_int_backend<>, boost::multiprecision::et_off>;

class SRP {
private:
  Int N;
  Int g;
  Int k;

public:
  SRP(Int N, Int g, Int k) : N(N), g(g), k(k) {};

  std::string big_int_to_hex(Int n, bool ignore = false) {
    std::stringstream ss;
    ss << std::hex << n;
    std::string str = ss.str();
    if (!ignore && str.length() % 2 == 1) {
      str.insert(0, "0");
    }

    std::transform(str.begin(), str.end(), str.begin(), ::tolower);
    return str;
  }

  Int hex_to_big_int(std::string hex) {
    if (hex.substr(0, 2) == "0x") {
      return Int(hex);
    } else {
      return Int("0x" + hex);
    }
  }

  bool is_zero_mod_n(std::string number) {
    Int n = hex_to_big_int(number);
    return int(n % N) == 0;
  }

  template<typename T, typename... Args>
  std::string H(std::initializer_list<T> list) {
    std::stringstream data;
    int i = 0;
    for (auto elem : list)
      i++ != 0 ? data << ":" << elem : data << elem;
    return sw::sha512::calculate(data.str());
  }

  Int random_big_number_512bit() {
    std::random_device rd;
    mt19937 base_gen(rd());
    independent_bits_engine<mt11213b, 512, Int> gen(base_gen);
    return gen();
  }

  std::string rnd() {
    return big_int_to_hex(random_big_number_512bit());
  }

  std::string s() {
    return rnd();
  }

  std::string x(std::string s, std::string p) {
    return H({s, p});
  }

  std::string v(std::string x) {
    return big_int_to_hex(powm(g, hex_to_big_int(x), N));
  }

  std::string a() {
    return rnd();
  }

  std::string A(std::string _a) {
    Int a = hex_to_big_int(_a);
    return big_int_to_hex(powm(g, a, N));
  }

  std::string b() {
    return rnd();
  }

  std::string B(std::string _v, std::string _b) {
    Int v = hex_to_big_int(_v);
    Int b = hex_to_big_int(_b);
    Int step1 = powm(g, b, N);
    return big_int_to_hex(k * v + step1);
  }

  std::string u(std::string A, std::string B) {
    return H({A, B});
  }

  std::string server_s(std::string _A, std::string _v, std::string _u, std::string _b) {
    Int A = hex_to_big_int(_A);
    Int v = hex_to_big_int(_v);
    Int b = hex_to_big_int(_b);
    Int u = hex_to_big_int(_u);
    Int step1 = powm(v, u, N);
    Int step2 = A * step1;
    Int final_result = powm(step2, b, N);
    return big_int_to_hex(final_result, true);
  }

  std::string client_s(std::string _B, std::string _x, std::string _u, std::string _a) {
    Int B = hex_to_big_int(_B);
    Int x = hex_to_big_int(_x);
    Int a = hex_to_big_int(_a);
    Int u = hex_to_big_int(_u);
    Int step1 = powm(g, x, N);
    Int step2 = k * step1;
    Int step3 = B - step2;
    Int step4 = u * x;
    Int step5 = a + step4;
    Int final_result = powm(step3, step5, N);
    return big_int_to_hex(final_result, true);
  }

  std::string M(std::string I, std::string s, std::string A, std::string B, std::string K) {
    std::string N_hex = big_int_to_hex(N);
    std::string g_hex = big_int_to_hex(g);
    Int N_hash_int = hex_to_big_int(H({N_hex}));
    Int g_hash_int = hex_to_big_int(H({g_hex}));
    std::string I_hash = H({I});
    std::string xor_step = big_int_to_hex(N_hash_int xor g_hash_int);

    return H({xor_step, I_hash, s, A, B, K});
  }

  std::string R(std::string A, std::string M, std::string K) {
    return H({A, M, K});
  }
};