/*

Copyright (C) 2015, 2016 Olaf Till <i7tiol@t-online.de>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; If not, see <http://www.gnu.org/licenses/>.

*/

#ifndef __OCT_PARALLEL_SENSITIVE_DATA__

#define __OCT_PARALLEL_SENSITIVE_DATA__

#include <string.h>

class
sensitive_base
{
public:
  virtual void zero_val (void) = 0;
};

class
zero_sensitive_base
{
public:
  zero_sensitive_base (void) {}

  ~zero_sensitive_base (void)
  {
    for (std::list<sensitive_base *>::const_iterator it = list.begin ();
         it != list.end (); it++)
      (*it)->zero_val ();
  }

  void add (sensitive_base *el)
  {
    list.push_front (el);
  }

private:
  std::list<sensitive_base *> list;
};

class
verifier: public sensitive_base
{
public:

  verifier (void)
    : valid (false), locked_length (0)
    {
      val.data = NULL;

      val.size = 0;
    }

  verifier (const char *username, const char *passwd,
            const gnutls_datum_t *salt, bool lock = false)
    : valid (false), locked_length (0)
    {
      val.data = NULL;

      val.size = 0;

      do_assign (username, passwd, salt, lock);
    }

  void assign (const char *username, const char *passwd,
               const gnutls_datum_t *salt, bool lock = false)
    {
      do_assign (username, passwd, salt, lock);
    }

  void free_val (void)
  {
    if (val.data)
      {
        zero_val ();

        if (locked_length)
          {
            if (munlock (val.data, locked_length))
              fprintf (stderr, "Could not unlock memory pages!\n");
            else
              locked_length = 0;
          }

        free (val.data);

        val.data = NULL;
      }

    val.size = 0;

    valid = false;
  }

  void zero_val (void)
  {
    valid = false;

    if (val.data)
      memset ((void *) val.data, '\0', val.size);
  }

  ~verifier (void)
  {
    free_val ();
  }

  bool good (void) {return valid;}

  const gnutls_datum_t &get_val (void) {return val;}

private:

  void do_assign (const char *&username, const char *&passwd,
                  const gnutls_datum_t *&salt, bool lock)
    {
      if (gnutls_srp_verifier (username, passwd, salt,
                               &gnutls_srp_2048_group_generator,
                               &gnutls_srp_2048_group_prime,
                               &val) == GNUTLS_E_SUCCESS &&
          val.data)
        {
          if (lock)
            {
              if (mlock (val.data, val.size))
                c_verror ("could not lock memory pages");
              else
                {
                  valid = true;

                  locked_length = val.size;
                }
            }
          else
            valid = true;
        }
    }

  gnutls_datum_t val;

  bool valid;

  int locked_length;
};

class
sensitive_string: public sensitive_base
{
public:

  sensitive_string (void)
  : valid (false), internally_allocated (false), locked_length (0)
    {
      val.data = NULL;

      val.size = 0;
    }

  sensitive_string (char *s, bool lock = false)
  : valid (false), internally_allocated (false), locked_length (0)
    {
      assign (s, lock);
    }

  sensitive_string (int n, bool lock = false)
  : valid (false), internally_allocated (false), locked_length (0)
    {
      val.data = NULL;

      val.size = 0;

      fill (n, lock);
    }

  void fill (int n, bool lock = false);

  void assign (char *s, bool lock = false)
    {
      if (valid || locked_length || internally_allocated)
        {
          valid = false;

          return;
        }

      val.data = (unsigned char *) s;

      val.size = strlen ((char *) val.data);

      valid = true;

     if (val.size != 0 && lock)
       {
         if (mlock (val.data, val.size))
           {
             c_verror ("could not lock memory pages");

             valid = false;
           }
         else
           locked_length = val.size;
       }
    }

  void free_val (void)
  {
    if (val.data)
      {
        zero_val ();

        if (locked_length)
          {
            if (munlock (val.data, locked_length))
              fprintf (stderr, "Could not unlock memory pages!\n");
            else
              locked_length = 0;
          }

        if (internally_allocated)
          {
            delete [] val.data;

            internally_allocated = false;
          }

        val.data = NULL;
      }
    val.size = 0;

    valid = false;
  }

  void zero_val (void)
  {
    valid = false;

    if (val.data)
      memset ((void *) val.data, '\0', val.size);
  }

  ~sensitive_string (void)
  {
    free_val ();
  }

  bool good (void) {return valid;}

  const gnutls_datum_t &get_val (void) {return val;}

  char *get_string (void) {return (char *) val.data;}

  int get_size (void) {return val.size;}

private:

  // val.size is the length of the string (without terminating '\0'),
  // not the length of the data. This is because gnutls encoding
  // functions expect a gnutls_datum_t with binary data.
  gnutls_datum_t val;

  bool valid;

  bool internally_allocated;

  int locked_length;
};

class
srp_base64_encode: public sensitive_base
{
public:

  srp_base64_encode (void)
  : valid (false), locked_length (0)
  {
    val.data = NULL;

    val.size = 0;
  }

  srp_base64_encode (const gnutls_datum_t *data, bool lock = false)
  : valid (false), locked_length (0)
  {
    val.data = NULL;

    val.size = 0;

    encode (data, lock);
  }

  void encode (const gnutls_datum_t *data, bool lock = false)
  {
    if (gnutls_srp_base64_encode_alloc (data, &val))
      {
        c_verror ("base64 encoding failed");

        valid = false;
      }
    else
      {
        valid = true;

        if (lock)
          {
            if (mlock (val.data, val.size))
              {
                c_verror ("could not lock memory pages");

                valid = false;
              }
            else
              locked_length = val.size;
          }
      }
  }

  void free_val (void)
  {
    if (val.data)
      {
        zero_val ();

        if (locked_length)
          {
            if (munlock (val.data, locked_length))
              fprintf (stderr, "Could not unlock memory pages!\n");
            else
              locked_length = 0;
          }

        gnutls_free ((void *) val.data);

        val.data = NULL;
      }
    val.size = 0;
  }

  void zero_val (void)
  {
    valid = false;

    if (val.data)
      memset ((void *) val.data, '\0', val.size);
  }

  ~srp_base64_encode (void)
  {
    free_val ();
  }

  bool good (void) {return valid;}

  const gnutls_datum_t &get_val (void) {return val;}

  char *get_string (void) {return (char *) val.data;}

  int get_size (void) {return val.size;}

private:

  gnutls_datum_t val;

  bool valid;

  int locked_length;
};

#endif // __OCT_PARALLEL_SENSITIVE_DATA__
