/* SSL socket module

   SSL support based on patches by Brian E Gallew and Laszlo Kovacs.

   Re-worked a bit by Bill Janssen to add server-side support and

   certificate decoding.  Chris Stawarz contributed some non-blocking

   patches.

   This module is imported by ssl.py. It should *not* be used

   directly.

   XXX should partial writes be enabled, SSL_MODE_ENABLE_PARTIAL_WRITE?

   XXX what about SSL_MODE_AUTO_RETRY?

*/

#include "Python.h"

#ifdef WITH_THREAD

#include "pythread.h"

#define PySSL_BEGIN_ALLOW_THREADS { \

			PyThreadState *_save = NULL;  \

			if (_ssl_locks_count>0) {_save = PyEval_SaveThread();}

#define PySSL_BLOCK_THREADS	if (_ssl_locks_count>0){PyEval_RestoreThread(_save)};

#define PySSL_UNBLOCK_THREADS	if (_ssl_locks_count>0){_save = PyEval_SaveThread()};

#define PySSL_END_ALLOW_THREADS	if (_ssl_locks_count>0){PyEval_RestoreThread(_save);} \

		 }

#else	/* no WITH_THREAD */

#define PySSL_BEGIN_ALLOW_THREADS

#define PySSL_BLOCK_THREADS

#define PySSL_UNBLOCK_THREADS

#define PySSL_END_ALLOW_THREADS

#endif

enum py_ssl_error {

	/* these mirror ssl.h */

	PY_SSL_ERROR_NONE,

	PY_SSL_ERROR_SSL,

	PY_SSL_ERROR_WANT_READ,

	PY_SSL_ERROR_WANT_WRITE,

	PY_SSL_ERROR_WANT_X509_LOOKUP,

	PY_SSL_ERROR_SYSCALL,     /* look at error stack/return value/errno */

	PY_SSL_ERROR_ZERO_RETURN,

	PY_SSL_ERROR_WANT_CONNECT,

	/* start of non ssl.h errorcodes */

	PY_SSL_ERROR_EOF,         /* special case of SSL_ERROR_SYSCALL */

	PY_SSL_ERROR_INVALID_ERROR_CODE

};

enum py_ssl_server_or_client {

	PY_SSL_CLIENT,

	PY_SSL_SERVER

};

enum py_ssl_cert_requirements {

	PY_SSL_CERT_NONE,

	PY_SSL_CERT_OPTIONAL,

	PY_SSL_CERT_REQUIRED

};

enum py_ssl_version {

	PY_SSL_VERSION_SSL2,

	PY_SSL_VERSION_SSL3,

	PY_SSL_VERSION_SSL23,

	PY_SSL_VERSION_TLS1,

};

/* Include symbols from _socket module */

#include "socketmodule.h"

#if defined(HAVE_POLL_H)

#include <poll.h>

#elif defined(HAVE_SYS_POLL_H)

#include <sys/poll.h>

#endif

/* Include OpenSSL header files */

#include "openssl/rsa.h"

#include "openssl/crypto.h"

#include "openssl/x509.h"

#include "openssl/x509v3.h"

#include "openssl/pem.h"

#include "openssl/ssl.h"

#include "openssl/err.h"

#include "openssl/rand.h"

/* SSL error object */

static PyObject *PySSLErrorObject;

#ifdef WITH_THREAD

/* serves as a flag to see whether we've initialized the SSL thread support. */

/* 0 means no, greater than 0 means yes */

static unsigned int _ssl_locks_count = 0;

#endif /* def WITH_THREAD */

/* SSL socket object */

#define X509_NAME_MAXLEN 256

/* RAND_* APIs got added to OpenSSL in 0.9.5 */

#if OPENSSL_VERSION_NUMBER >= 0x0090500fL

# define HAVE_OPENSSL_RAND 1

#else

# undef HAVE_OPENSSL_RAND

#endif

typedef struct {

	PyObject_HEAD

	PySocketSockObject *Socket;	/* Socket on which we're layered */

	SSL_CTX*	ctx;

	SSL*		ssl;

	X509*		peer_cert;

	char		server[X509_NAME_MAXLEN];

	char		issuer[X509_NAME_MAXLEN];

} PySSLObject;

static PyTypeObject PySSL_Type;

static PyObject *PySSL_SSLwrite(PySSLObject *self, PyObject *args);

static PyObject *PySSL_SSLread(PySSLObject *self, PyObject *args);

static int check_socket_and_wait_for_timeout(PySocketSockObject *s,

					     int writing);

static PyObject *PySSL_peercert(PySSLObject *self, PyObject *args);

static PyObject *PySSL_cipher(PySSLObject *self);

#define PySSLObject_Check(v)	(Py_TYPE(v) == &PySSL_Type)

typedef enum {

	SOCKET_IS_NONBLOCKING,

	SOCKET_IS_BLOCKING,

	SOCKET_HAS_TIMED_OUT,

	SOCKET_HAS_BEEN_CLOSED,

	SOCKET_TOO_LARGE_FOR_SELECT,

	SOCKET_OPERATION_OK

} timeout_state;

/* Wrap error strings with filename and line # */

#define STRINGIFY1(x) #x

#define STRINGIFY2(x) STRINGIFY1(x)

#define ERRSTR1(x,y,z) (x ":" y ": " z)

#define ERRSTR(x) ERRSTR1("_ssl.c", STRINGIFY2(__LINE__), x)

/* XXX It might be helpful to augment the error message generated

   below with the name of the SSL function that generated the error.

   I expect it's obvious most of the time.

*/

static PyObject *

PySSL_SetError(PySSLObject *obj, int ret, char *filename, int lineno)

{

	PyObject *v;

	char buf[2048];

	char *errstr;

	int err;

	enum py_ssl_error p = PY_SSL_ERROR_NONE;

	assert(ret <= 0);

	if (obj->ssl != NULL) {

		err = SSL_get_error(obj->ssl, ret);

		switch (err) {

		case SSL_ERROR_ZERO_RETURN:

			errstr = "TLS/SSL connection has been closed";

			p = PY_SSL_ERROR_ZERO_RETURN;

			break;

		case SSL_ERROR_WANT_READ:

			errstr = "The operation did not complete (read)";

			p = PY_SSL_ERROR_WANT_READ;

			break;

		case SSL_ERROR_WANT_WRITE:

			p = PY_SSL_ERROR_WANT_WRITE;

			errstr = "The operation did not complete (write)";

			break;

		case SSL_ERROR_WANT_X509_LOOKUP:

			p = PY_SSL_ERROR_WANT_X509_LOOKUP;

			errstr =

                            "The operation did not complete (X509 lookup)";

			break;

		case SSL_ERROR_WANT_CONNECT:

			p = PY_SSL_ERROR_WANT_CONNECT;

			errstr = "The operation did not complete (connect)";

			break;

		case SSL_ERROR_SYSCALL:

		{

			unsigned long e = ERR_get_error();

			if (e == 0) {

				if (ret == 0 || !obj->Socket) {

				  p = PY_SSL_ERROR_EOF;

				  errstr =

                                      "EOF occurred in violation of protocol";

				} else if (ret == -1) {

				  /* underlying BIO reported an I/O error */

                                  return obj->Socket->errorhandler();

				} else { /* possible? */

                                  p = PY_SSL_ERROR_SYSCALL;

                                  errstr = "Some I/O error occurred";

				}

			} else {

				p = PY_SSL_ERROR_SYSCALL;

				/* XXX Protected by global interpreter lock */

				errstr = ERR_error_string(e, NULL);

			}

			break;

		}

		case SSL_ERROR_SSL:

		{

			unsigned long e = ERR_get_error();

			p = PY_SSL_ERROR_SSL;

			if (e != 0)

				/* XXX Protected by global interpreter lock */

				errstr = ERR_error_string(e, NULL);

			else {	/* possible? */

				errstr =

                                    "A failure in the SSL library occurred";

			}

			break;

		}

		default:

			p = PY_SSL_ERROR_INVALID_ERROR_CODE;

			errstr = "Invalid error code";

		}

	} else {

		errstr = ERR_error_string(ERR_peek_last_error(), NULL);

	}

	PyOS_snprintf(buf, sizeof(buf), "_ssl.c:%d: %s", lineno, errstr);

	v = Py_BuildValue("(is)", p, buf);

	if (v != NULL) {

		PyErr_SetObject(PySSLErrorObject, v);

		Py_DECREF(v);

	}

	return NULL;

}

static PyObject *

_setSSLError (char *errstr, int errcode, char *filename, int lineno) {

	char buf[2048];

	PyObject *v;

	if (errstr == NULL) {

		errcode = ERR_peek_last_error();

		errstr = ERR_error_string(errcode, NULL);

	}

	PyOS_snprintf(buf, sizeof(buf), "_ssl.c:%d: %s", lineno, errstr);

	v = Py_BuildValue("(is)", errcode, buf);

	if (v != NULL) {

		PyErr_SetObject(PySSLErrorObject, v);

		Py_DECREF(v);

	}

	return NULL;

}

static PySSLObject *

newPySSLObject(PySocketSockObject *Sock, char *key_file, char *cert_file,

	       enum py_ssl_server_or_client socket_type,

	       enum py_ssl_cert_requirements certreq,

	       enum py_ssl_version proto_version,

	       char *cacerts_file)

{

	PySSLObject *self;

	char *errstr = NULL;

	int ret;

	int verification_mode;

	self = PyObject_New(PySSLObject, &PySSL_Type); /* Create new object */

	if (self == NULL)

		return NULL;

	memset(self->server, '\0', sizeof(char) * X509_NAME_MAXLEN);

	memset(self->issuer, '\0', sizeof(char) * X509_NAME_MAXLEN);

	self->peer_cert = NULL;

	self->ssl = NULL;

	self->ctx = NULL;

	self->Socket = NULL;

	/* Make sure the SSL error state is initialized */

	(void) ERR_get_state();

	ERR_clear_error();

	if ((key_file && !cert_file) || (!key_file && cert_file)) {

		errstr = ERRSTR("Both the key & certificate files "

                                "must be specified");

		goto fail;

	}

	if ((socket_type == PY_SSL_SERVER) &&

	    ((key_file == NULL) || (cert_file == NULL))) {

		errstr = ERRSTR("Both the key & certificate files "

                                "must be specified for server-side operation");

		goto fail;

	}

	PySSL_BEGIN_ALLOW_THREADS

	if (proto_version == PY_SSL_VERSION_TLS1)

		self->ctx = SSL_CTX_new(TLSv1_method()); /* Set up context */

	else if (proto_version == PY_SSL_VERSION_SSL3)

		self->ctx = SSL_CTX_new(SSLv3_method()); /* Set up context */

	else if (proto_version == PY_SSL_VERSION_SSL2)

		self->ctx = SSL_CTX_new(SSLv2_method()); /* Set up context */

	else if (proto_version == PY_SSL_VERSION_SSL23)

		self->ctx = SSL_CTX_new(SSLv23_method()); /* Set up context */

	PySSL_END_ALLOW_THREADS

	if (self->ctx == NULL) {

		errstr = ERRSTR("Invalid SSL protocol variant specified.");

		goto fail;

	}

	if (certreq != PY_SSL_CERT_NONE) {

		if (cacerts_file == NULL) {

			errstr = ERRSTR("No root certificates specified for "

                                  "verification of other-side certificates.");

			goto fail;

		} else {

			PySSL_BEGIN_ALLOW_THREADS

			ret = SSL_CTX_load_verify_locations(self->ctx,

							    cacerts_file,

                                                            NULL);

			PySSL_END_ALLOW_THREADS

			if (ret != 1) {

				_setSSLError(NULL, 0, __FILE__, __LINE__);

				goto fail;

			}

		}

	}

	if (key_file) {

		PySSL_BEGIN_ALLOW_THREADS

		ret = SSL_CTX_use_PrivateKey_file(self->ctx, key_file,

						  SSL_FILETYPE_PEM);

		PySSL_END_ALLOW_THREADS

		if (ret != 1) {

			_setSSLError(NULL, ret, __FILE__, __LINE__);

			goto fail;

		}

		PySSL_BEGIN_ALLOW_THREADS

		ret = SSL_CTX_use_certificate_chain_file(self->ctx,

							 cert_file);

		PySSL_END_ALLOW_THREADS

		if (ret != 1) {

			/*

			fprintf(stderr, "ret is %d, errcode is %lu, %lu, with file \"%s\"\n",

				ret, ERR_peek_error(), ERR_peek_last_error(), cert_file);

				*/

			if (ERR_peek_last_error() != 0) {

				_setSSLError(NULL, ret, __FILE__, __LINE__);

				goto fail;

			}

		}

	}

        /* ssl compatibility */

        SSL_CTX_set_options(self->ctx, SSL_OP_ALL);

	verification_mode = SSL_VERIFY_NONE;

	if (certreq == PY_SSL_CERT_OPTIONAL)

		verification_mode = SSL_VERIFY_PEER;

	else if (certreq == PY_SSL_CERT_REQUIRED)

		verification_mode = (SSL_VERIFY_PEER |

				     SSL_VERIFY_FAIL_IF_NO_PEER_CERT);

	SSL_CTX_set_verify(self->ctx, verification_mode,

			   NULL); /* set verify lvl */

	PySSL_BEGIN_ALLOW_THREADS

	self->ssl = SSL_new(self->ctx); /* New ssl struct */

	PySSL_END_ALLOW_THREADS

	SSL_set_fd(self->ssl, Sock->sock_fd);	/* Set the socket for SSL */

	/* If the socket is in non-blocking mode or timeout mode, set the BIO

	 * to non-blocking mode (blocking is the default)

	 */

	if (Sock->sock_timeout >= 0.0) {

		/* Set both the read and write BIO's to non-blocking mode */

		BIO_set_nbio(SSL_get_rbio(self->ssl), 1);

		BIO_set_nbio(SSL_get_wbio(self->ssl), 1);

	}

	PySSL_BEGIN_ALLOW_THREADS

	if (socket_type == PY_SSL_CLIENT)

		SSL_set_connect_state(self->ssl);

	else

		SSL_set_accept_state(self->ssl);

	PySSL_END_ALLOW_THREADS

	self->Socket = Sock;

	Py_INCREF(self->Socket);

	return self;

 fail:

	if (errstr)

		PyErr_SetString(PySSLErrorObject, errstr);

	Py_DECREF(self);

	return NULL;

}

static PyObject *

PySSL_sslwrap(PyObject *self, PyObject *args)

{

	PySocketSockObject *Sock;

	int server_side = 0;

	int verification_mode = PY_SSL_CERT_NONE;

	int protocol = PY_SSL_VERSION_SSL23;

	char *key_file = NULL;

	char *cert_file = NULL;

	char *cacerts_file = NULL;

	if (!PyArg_ParseTuple(args, "O!i|zziiz:sslwrap",

			      PySocketModule.Sock_Type,

			      &Sock,

			      &server_side,

			      &key_file, &cert_file,

			      &verification_mode, &protocol,

			      &cacerts_file))

		return NULL;

	/*

	fprintf(stderr,

		"server_side is %d, keyfile %p, certfile %p, verify_mode %d, "

		"protocol %d, certs %p\n",

		server_side, key_file, cert_file, verification_mode,

		protocol, cacerts_file);

	 */

	return (PyObject *) newPySSLObject(Sock, key_file, cert_file,

					   server_side, verification_mode,

					   protocol, cacerts_file);

}

PyDoc_STRVAR(ssl_doc,

"sslwrap(socket, server_side, [keyfile, certfile, certs_mode, protocol,\n"

"                              cacertsfile]) -> sslobject");

/* SSL object methods */

static PyObject *PySSL_SSLdo_handshake(PySSLObject *self)

{

	int ret;

	int err;

	int sockstate;

	/* Actually negotiate SSL connection */

	/* XXX If SSL_do_handshake() returns 0, it's also a failure. */

	sockstate = 0;

	do {

		PySSL_BEGIN_ALLOW_THREADS

		ret = SSL_do_handshake(self->ssl);

		err = SSL_get_error(self->ssl, ret);

		PySSL_END_ALLOW_THREADS

		if(PyErr_CheckSignals()) {

			return NULL;

		}

		if (err == SSL_ERROR_WANT_READ) {

			sockstate = check_socket_and_wait_for_timeout(self->Socket, 0);

		} else if (err == SSL_ERROR_WANT_WRITE) {

			sockstate = check_socket_and_wait_for_timeout(self->Socket, 1);

		} else {

			sockstate = SOCKET_OPERATION_OK;

		}

		if (sockstate == SOCKET_HAS_TIMED_OUT) {

			PyErr_SetString(PySSLErrorObject,

				ERRSTR("The handshake operation timed out"));

			return NULL;

		} else if (sockstate == SOCKET_HAS_BEEN_CLOSED) {

			PyErr_SetString(PySSLErrorObject,

				ERRSTR("Underlying socket has been closed."));

			return NULL;

		} else if (sockstate == SOCKET_TOO_LARGE_FOR_SELECT) {

			PyErr_SetString(PySSLErrorObject,

			  ERRSTR("Underlying socket too large for select()."));

			return NULL;

		} else if (sockstate == SOCKET_IS_NONBLOCKING) {

			break;

		}

	} while (err == SSL_ERROR_WANT_READ || err == SSL_ERROR_WANT_WRITE);

	if (ret < 1)

		return PySSL_SetError(self, ret, __FILE__, __LINE__);

	self->ssl->debug = 1;

	if (self->peer_cert)

		X509_free (self->peer_cert);

	PySSL_BEGIN_ALLOW_THREADS

	if ((self->peer_cert = SSL_get_peer_certificate(self->ssl))) {

		X509_NAME_oneline(X509_get_subject_name(self->peer_cert),

				  self->server, X509_NAME_MAXLEN);

		X509_NAME_oneline(X509_get_issuer_name(self->peer_cert),

				  self->issuer, X509_NAME_MAXLEN);

	}

	PySSL_END_ALLOW_THREADS

	Py_INCREF(Py_None);

	return Py_None;

}

static PyObject *

PySSL_server(PySSLObject *self)

{

	return PyString_FromString(self->server);

}

static PyObject *

PySSL_issuer(PySSLObject *self)

{

	return PyString_FromString(self->issuer);

}

static PyObject *

_create_tuple_for_attribute (ASN1_OBJECT *name, ASN1_STRING *value) {

	char namebuf[X509_NAME_MAXLEN];

	int buflen;

	PyObject *name_obj;

	PyObject *value_obj;

	PyObject *attr;

	unsigned char *valuebuf = NULL;

	buflen = OBJ_obj2txt(namebuf, sizeof(namebuf), name, 0);

	if (buflen < 0) {

		_setSSLError(NULL, 0, __FILE__, __LINE__);

		goto fail;

	}

	name_obj = PyString_FromStringAndSize(namebuf, buflen);

	if (name_obj == NULL)

		goto fail;

	buflen = ASN1_STRING_to_UTF8(&valuebuf, value);