// python interface to openssl BIGNUM prime and modular inverse routines
//
//   bignum.generate_prime() -> BN_generate_prime_ex()
//
//   bignum.is_prime() -> BN_is_prime_ex()
//
//   bignum.mod_inverse() -> BN_mod_inverse()
//
// R. Perry, May 2017; updated Feb. 2019
//
// References:
//
//   https://docs.python.org/3/extending/extending.html
//   https://docs.python.org/3/c-api/
//   https://www.openssl.org/docs/man1.1.1/man3/BN_generate_prime.html
//   https://www.openssl.org/docs/man1.1.1/man3/BN_mod_inverse.html

#include "Python.h"
#include <openssl/bn.h>
#include <openssl/err.h> // void ERR_print_errors_fp(FILE *fp);

//------------------------------------------------------------------------------
//
// error-checking interfaces to BN_new(void), BN_CTX_new(), BN_GENCB_new(), BN_bin2bn()

// BIGNUM *BN_new(void); // allocate and initialize BIGNUM
//
//   BN_new() allocates and initializes a BIGNUM structure.
//   If the allocation fails, returns NULL.
//
static BIGNUM *bn_new( void)
{
  BIGNUM *r = BN_new();

  if( !r) { ERR_print_errors_fp( stdout); PyErr_SetString( PyExc_RuntimeError, "BN_new() failed"); }

  return r;
}

// BN_CTX *BN_CTX_new(void); // allocate and initialize BN_CTX structure
// 
//   BN_CTX_new() allocates and initializes a BN_CTX structure.
//   If the allocation fails, returns NULL.
// 
static BN_CTX *bn_ctx_new( void)
{
  BN_CTX *r = BN_CTX_new();

  if( !r) { ERR_print_errors_fp( stdout); PyErr_SetString( PyExc_RuntimeError, "BN_CTX_new() failed"); }

  return r;
}

// BN_GENCB *BN_GENCB_new(void); // create BN_GENCB structure
//
//   BN_GENCB_new() returns a pointer to a BN_GENCB structure on success, or NULL otherwise.
//
static BN_GENCB *bn_gencb_new( void)
{
  BN_GENCB *r = BN_GENCB_new();

  if( !r) { ERR_print_errors_fp( stdout); PyErr_SetString( PyExc_RuntimeError, "BN_GENCB_new() failed"); }

  return r;
}

// BIGNUM *BN_bin2bn(const unsigned char *s, int len, BIGNUM *ret); // convert byte array to BIGNUM
//
//   BN_bin2bn() converts the positive integer in big-endian form of length len
//   at s into a BIGNUM and places it in ret. If ret is NULL, a new BIGNUM is created.
//   BN_bin2bn() returns the BIGNUM, NULL on error.
//
static BIGNUM *bn_bin2bn( const unsigned char *s, size_t len, BIGNUM *ret)
{
  BIGNUM *r = BN_bin2bn( s, len, ret);

  if( !r) { ERR_print_errors_fp( stdout); PyErr_SetString( PyExc_RuntimeError, "BN_bin2bn() failed"); }

  return r;
}

//------------------------------------------------------------------------------
//
// callback functions should return 1 to continue, 0 to stop
// returning 0 will make it crash, e.g.: RuntimeError: BN_is_prime_ex() failed
//
// The callback interface is not very well documented, and only seems useful
// to indicate progress.  It is used in apps/ dsaparam.c, gendh.c, dhparam.c,
// and genrsa.c, similar to this usage from crypto/bn/bntest.c
//
//   static int genprime_cb(int p, int n, BN_GENCB *arg) { char c='*';
//     if (p == 0) c='.'; if (p == 1) c='+'; if (p == 2) c='*'; if (p == 3) c='\n';
//     putc(c, stderr); fflush(stderr); return 1; }
//
// man BN_GENCB_call:
//
//   int BN_GENCB_call(BN_GENCB *cb, int a, int b)
//
//   BN_GENCB_call calls the callback function held in the BN_GENCB structure
//   and passes the ints a and b as arguments.
//
//   a BN_GENCB structure should be initialised with a call to BN_GENCB_set,
//   where gencb is a BN_GENCB *, callback is of type
//   int (*callback)(int, int, BN_GENCB *) and cb_arg is a void *.
//
// bn.h:
//
//   #define BN_GENCB_set(gencb, callback, cb_arg) ...

typedef struct { PyObject *cb, *cb_arg; } CBS;

static int callback( int a, int b, BN_GENCB *cb)
{
  CBS *cbs = BN_GENCB_get_arg( cb);

#if 0 // debug
  printf( "callback: %i %i %p %p %p\n", a, b, cbs, cbs->cb, cbs->cb_arg);
#endif

 // construct Python callback argument list
 //
  PyObject *arglist = Py_BuildValue( "(iiO)", a, b, cbs->cb_arg);

#if 0 // debug
  printf( "callback: arglist = %p\n", arglist);
#endif

  if( arglist == NULL) return 0;

 // call the Python callback function
 //
  PyObject *result = PyObject_CallObject( cbs->cb, arglist);

#if 0 // debug
  printf( "callback: result = %p\n", result);
#endif

  Py_DECREF( arglist); // done with arglist

  if( result == NULL) return 0;

  int r = PyLong_AsLong( result);

  Py_DECREF( result); // done with result

  return r;
}
 
//------------------------------------------------------------------------------

static PyObject *bignum_generate_prime( PyObject *self, PyObject *args, PyObject *keywords)
{
  int nbits, safe = 0;  PyObject *add = NULL, *rem = NULL, *cb = NULL, *cb_arg = Py_None;

  static char *keywordlist[] = { "nbits", "safe", "add", "rem", "cb", "cb_arg", NULL };

  if( !PyArg_ParseTupleAndKeywords( args, keywords, "i|pO!O!OO", keywordlist,
      &nbits, &safe, &PyLong_Type, &add, &PyLong_Type, &rem, &cb, &cb_arg) ) return NULL;

  if( cb && !PyCallable_Check( cb) )
    { PyErr_SetString( PyExc_TypeError, "callback argument cb must be callable"); return NULL; }

#if 0 // debug
  printf( "safe = %i\n", safe);
#endif

 // int BN_generate_prime_ex( BIGNUM *ret, int bits, int safe, const BIGNUM *add,
 //   const BIGNUM *rem, BN_GENCB *cb)
 //
 // returns 1 on success or 0 on error.
 //
 // If safe is true, it will be a safe prime (i.e. a prime p so that (p-1)/2 is also prime).
 //
 // If add is not NULL, the prime will fulfill the condition p % add == rem
 // (p % add == 1 if rem == NULL) in order to suit a given generator.
 // 
 // If cb is not NULL, it is used as follows:
 //  - BN_GENCB_call(cb, 0, i) is called after generating the i-th potential prime number.
 //  - While the number is being tested for primality, BN_GENCB_call(cb, 1, j) is called
 //    as described below.
 //  - When a prime has been found, BN_GENCB_call(cb, 2, i) is called.
 //    [RP: only if safe is non-zero]

  BIGNUM *cadd = NULL, *crem = NULL;

 // handle optional add and rem arguments
 //
  if( add)
  {
    // see comments in bignum_is_prime() below regarding _PyLong_NumBits(), etc.

    size_t n = (_PyLong_NumBits(add) >> 3) + 1;  unsigned char b[n];

    if( _PyLong_AsByteArray( (PyLongObject *) add, b, n, 0, 0) ) return NULL;

    if( !(cadd = bn_bin2bn( b, n, NULL)) ) return NULL;
  }

  if( rem)
  {
    size_t n = (_PyLong_NumBits(rem) >> 3) + 1;  unsigned char b[n];

    if( _PyLong_AsByteArray( (PyLongObject *) rem, b, n, 0, 0) ) { BN_free( cadd); return NULL; }

    if( !(crem = bn_bin2bn( b, n, NULL)) ) { BN_free( cadd); return NULL; }
  }

 // handle optional callback argument
 //
  BN_GENCB *ccbp = NULL; CBS cbs = { cb, cb_arg }; 

  if( cb)
  {
    if( !(ccbp = bn_gencb_new()) ) { BN_free( cadd); BN_free( crem); return NULL; }

    BN_GENCB_set( ccbp, callback, &cbs);
  }

 // generate prime p using BN_generate_prime_ex()
 //
  BIGNUM *p = bn_new(); int r = 0;

  if( p) r = BN_generate_prime_ex( p, nbits, safe, cadd, crem, ccbp);

 // cleanup
 //
  BN_free( cadd); BN_free( crem); BN_GENCB_free( ccbp);

 // error check
 //
  if( !p) return NULL; // bn_new() error

  if( !r) { ERR_print_errors_fp( stdout);
    PyErr_SetString( PyExc_RuntimeError, "BN_generate_prime_ex() failed"); BN_free( p); return NULL; }

 // convert the result to a byte array
 // 
 // int BN_bn2bin( const BIGNUM *a, unsigned char *to)
 //
 //   BN_bn2bin() converts the absolute value of a into big-endian form
 //   and stores it at to. to must point to BN_num_bytes(a) bytes of memory.
 //
  int n = BN_num_bytes( p);  unsigned char b[n];  BN_bn2bin( p, b);  BN_free( p);

#if 0 // debug
  printf( "b = "); for( int i = 0; i < n; ++i) { printf( "%02x", b[i]); } putchar('\n');
#endif

 // convert the byte array to a Python Long
 //
 // undocumented, from Include/longobject.h, #ifndef Py_LIMITED_API
 //
 // PyAPI_FUNC(PyObject *) _PyLong_FromByteArray( const unsigned char* bytes, size_t n,
 //     int little_endian, int is_signed)
 //
 //   View the n unsigned bytes as a binary integer in base 256,
 //   and return a Python int with the same numeric value.
 //
  return _PyLong_FromByteArray( b, n, 0, 0);
}

//------------------------------------------------------------------------------

static PyObject *bignum_is_prime( PyObject *self, PyObject *args, PyObject *keywords)
{
  PyObject *p, *cb = NULL, *cb_arg = Py_None;  int nchecks = BN_prime_checks;

  static char *keywordlist[] = { "p", "nchecks", "cb", "cb_arg", NULL };

  if( !PyArg_ParseTupleAndKeywords( args, keywords, "O!|iOO", keywordlist, &PyLong_Type, &p, &nchecks, &cb, &cb_arg) )
    return NULL;

  if( cb && !PyCallable_Check( cb) )
    { PyErr_SetString( PyExc_TypeError, "callback argument cb must be callable"); return NULL; }

#if 0 // debug
  printf( "nchecks = %i\n", nchecks);
#endif

 // convert Python Long to a byte array
 //
 // undocumented, from Include/longobject.h, #ifndef Py_LIMITED_API
 //
 // PyAPI_FUNC(int) _PyLong_AsByteArray(PyLongObject* v, unsigned char* bytes, size_t n,
 //     int little_endian, int is_signed);
 //
 //   Convert the least-significant 8*n bits of long v to a base-256 integer,
 //   stored in array bytes.  Normally return 0, return -1 on error.
 //
 //   [RP: see example usage in Modules/_pickle.c]
 //
 // PyAPI_FUNC(int) _PyLong_Sign(PyObject *v)
 //
 //   Return 0 if v is 0, -1 if v < 0, +1 if v > 0.
 //
 // PyAPI_FUNC(size_t) _PyLong_NumBits(PyObject *v)
 //
 //   Return the number of bits needed to represent the absolute value of a long.
 //   For example, this returns 1 for 1 and -1, 2 for 2 and -2, and 2 for 3 and -3.
 //   It returns 0 for 0.
 //
  size_t n = (_PyLong_NumBits(p) >> 3) + 1;  unsigned char b[n];

  if( _PyLong_AsByteArray( (PyLongObject *) p, b, n, 0, 0) ) return NULL;

  BIGNUM *bn = bn_bin2bn( b, n, NULL);  if( !bn) return NULL;

  BN_CTX *ctx = bn_ctx_new();  if( !ctx) { BN_free( bn); return NULL; }

 // int BN_is_prime_ex( const BIGNUM *p, int nchecks, BN_CTX *ctx, BN_GENCB *cb)
 //
 //   returns 0 if the number is composite, 1 if it is prime with an error probability
 //   of less than 0.25^nchecks, and -1 on error.
 //
 //   If nchecks == BN_prime_checks, a number of iterations is used that yields
 //   a false positive rate of at most 2^-80 for random input.
 //
 //     [RP: bn.h defines BN_prime_checks as 0 and says:
 //     "select number of iterations based on the size of the number"]
 //
 //   If cb is not NULL, BN_GENCB_call(cb, 1, j) is called after the j-th iteration
 //   (j = 0, 1, ...).
 // 
 // handle optional callback argument
 //
  BN_GENCB *ccbp = NULL;  CBS cbs = { cb, cb_arg }; 

  if( cb)
  {
    if( !(ccbp = bn_gencb_new()) ) { BN_free( bn); BN_CTX_free( ctx); return NULL; }

    BN_GENCB_set( ccbp, callback, &cbs);
  }

 // call BN_is_prime_ex()
 //
  int r = BN_is_prime_ex( bn, nchecks, ctx, ccbp);

 // cleanup
 //
  BN_free( bn);  BN_CTX_free( ctx);  BN_GENCB_free(ccbp);

 // normal return
 //
  if( r == 1) Py_RETURN_TRUE; else if( r == 0) Py_RETURN_FALSE;

 // error return
 //
  ERR_print_errors_fp( stdout); PyErr_SetString( PyExc_RuntimeError, "BN_is_prime_ex() failed"); 

  return NULL;
}

//------------------------------------------------------------------------------

static PyObject *bignum_mod_inverse( PyObject *self, PyObject *args)
{
  PyObject *a, *n;

  if( !PyArg_ParseTuple( args, "O!O!", &PyLong_Type, &a, &PyLong_Type, &n) ) return NULL;

 // convert Python Long to byte array, then convert byte array to BIGNUM
 //
 // see comments in bignum_is_prime() regarding _PyLong_NumBits(), etc.
 //
  size_t na = (_PyLong_NumBits(a) >> 3) + 1;  unsigned char ba[na];

  size_t nn = (_PyLong_NumBits(n) >> 3) + 1;  unsigned char bn[nn];

  if( _PyLong_AsByteArray( (PyLongObject *) a, ba, na, 0, 0) ) return NULL;

  if( _PyLong_AsByteArray( (PyLongObject *) n, bn, nn, 0, 0) ) return NULL;

  BIGNUM *ca = bn_bin2bn( ba, na, NULL);  if( !ca) return NULL;

  BIGNUM *cn = bn_bin2bn( bn, nn, NULL);  if( !cn) { BN_free( ca); return NULL; }

  BN_CTX *ctx = bn_ctx_new();  if( !ctx) { BN_free( ca); BN_free( cn); return NULL; }

 // BIGNUM *BN_mod_inverse(BIGNUM *r, BIGNUM *a, const BIGNUM *n, BN_CTX *ctx)
 //
 //   BN_mod_inverse() computes the inverse of a modulo n places the result in r
 //   ("(a*r)%n==1"). If r is NULL, a new BIGNUM is created.
 //
 //   ctx is a previously allocated BN_CTX used for temporary variables.
 //   r may be the same BIGNUM as a or n.
 //
 //   BN_mod_inverse() returns the BIGNUM containing the inverse, and NULL on error.
 //
  BIGNUM *r = BN_mod_inverse( NULL, ca, cn, ctx);

 // cleanup
 //
  BN_free( ca);  BN_free( cn);  BN_CTX_free( ctx);

 // error check
 //
  if( !r) { ERR_print_errors_fp( stdout);
    PyErr_SetString( PyExc_RuntimeError, "BN_mod_inverse() failed"); return NULL; }

 // convert the result to a byte array, then convert the byte array to a Python Long
 // 
  int nr = BN_num_bytes( r);  unsigned char br[nr];  BN_bn2bin( r, br); BN_free( r);

  return _PyLong_FromByteArray( br, nr, 0, 0);
}

//------------------------------------------------------------------------------

PyDoc_STRVAR( module_doc, "openssl bignum prime routines");

static PyMethodDef bignum_methods[] =
{
  // The casts of bignum_generate_prime and is_prime are necessary since PyCFunction
  // values only take two PyObject* parameters, but bignum_generate_prime() and
  // isprime() take three because of the keyword list.
  //
  { "generate_prime", (PyCFunction) bignum_generate_prime, METH_VARARGS | METH_KEYWORDS,
	PyDoc_STR( "generate_prime( nbits, safe=False, add=None, rem=1, cb=None, cb_arg=None)\n\n\
Return a random prime.\n\n\
If safe is true, it will be a safe prime (i.e. a prime p so that (p-1)/2 is also prime).\n\n\
If add is specified, the prime will fulfill the condition p % add == rem.\n\n\
If the callback cb is specified it will be called after generating the i'th\n\
potential prime number with arguments (0,i,cb_arg); while the number is being\n\
tested for primality, it will be called as from is_prime(); if safe is true it\n\
will be called when a prime has been found with arguments (2,i,cb_arg)."
)},

  { "is_prime", (PyCFunction) bignum_is_prime, METH_VARARGS | METH_KEYWORDS,
	PyDoc_STR( "is_prime( p, nchecks=0, cb=None, cb_arg=None)\n\n\
Test if a number is prime.\n\n\
With nchecks=0 a number of iterations is used that yields a false positive rate\n\
of at most 2**-80 for random input.\n\n\
If the callback cb is specified it will be called after the j'th iteration\n\
with arguments (1,j,cb_arg)."
)},

  { "mod_inverse", bignum_mod_inverse, METH_VARARGS,
	PyDoc_STR( "mod_inverse( a, n)\n\n\
Returns the inverse of a modulo n."
)},

  { NULL, NULL}
};

static struct PyModuleDef bignum_module =
{
  PyModuleDef_HEAD_INIT,
  "bignum",
  module_doc,
  0,
  bignum_methods,
  NULL,
  NULL,
  NULL,
  NULL
};

PyMODINIT_FUNC PyInit_bignum( void)
{
  return PyModule_Create( &bignum_module);
}