py-pkstl 0.1.0

extern crate pkstl;

use std::io::BufWriter;
use pyo3::prelude::*;
use pyo3::types::*;
use pyo3::exceptions::*;
use ring::signature::KeyPair;

#[pyclass]
struct Ed25519KeyPair {
	inner: ring::signature::Ed25519KeyPair
}

#[pymethods]
impl Ed25519KeyPair {
	
	/// Instantiate key pair from seed
	/// Arguments:
	/// - seed: bytes (length=32)
	/// Returns: Ed255KeyPair
	#[staticmethod]
	fn from_seed_unchecked(seed: Py<PyBytes>) -> PyResult<Ed25519KeyPair> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		match ring::signature::Ed25519KeyPair::from_seed_unchecked(seed.as_ref(py).as_bytes()) {
			Ok(val) => Ok(Ed25519KeyPair { inner: val }),
			Err(_) => Err(ValueError::py_err("Error"))
		}
	}
	
	/// Public key
	#[getter]
	fn get_pubkey(&self) -> PyResult<PyObject> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		Ok(PyBytes::new(py, self.inner.public_key().as_ref()).to_object(py))
	}
}

#[pyclass]
struct Seed32 {
	inner: pkstl::Seed32
}

#[pymethods]
impl Seed32 {
	
	/// Create random seed
	/// Returns: Seed32
	#[staticmethod]
	fn random() -> PyResult<Seed32> {
		Ok(Seed32 { inner: pkstl::Seed32::random() })
	}
	
	/// Seed bytes
	#[getter]
	fn get_bytes(&self) -> PyResult<PyObject> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		Ok(PyBytes::new(py, self.inner.as_ref()).to_object(py))
	}
}

#[pyclass]
struct SecureLayerConfig {
	inner: pkstl::SecureLayerConfig
}

#[pymethods]
impl SecureLayerConfig {
	
	/// Instantiate SecureLayerConfig with default attributes
	/// Returns: SecureLayerConfig
	#[staticmethod]
	fn default() -> PyResult<SecureLayerConfig> {
		Ok(SecureLayerConfig { inner: pkstl::SecureLayerConfig::default() })
	}
}

#[pyclass]
struct MinimalSecureLayer {
	inner: pkstl::MinimalSecureLayer
}

#[pymethods]
impl MinimalSecureLayer {
	
	/// Instantiate MinimalSecureLayer
	/// Arguments:
	/// - config: SecureLayerConfig
	/// - expected_remote_sig_public_key: bytes (optional)
	/// Returns: MinimalSecureLayer
	#[staticmethod]
	fn create(config: &SecureLayerConfig, expected_remote_sig_public_key: Option<Py<PyBytes>>) -> PyResult<MinimalSecureLayer> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		match pkstl::MinimalSecureLayer::create(
			config.inner,
			match expected_remote_sig_public_key {
				Some(val) => Some(val.as_ref(py).as_bytes().to_vec()),
				None => None
			}
		) {
			Ok(minimal_secure_layer) => Ok(MinimalSecureLayer { inner: minimal_secure_layer }),
			Err(_) => Err(ValueError::py_err("Error"))
		}
	}
}

#[pyclass]
struct IncomingBinaryMessage {
	inner: pkstl::IncomingBinaryMessage
}

#[pymethods]
impl IncomingBinaryMessage {
	
	/// Message type
	/// 1=Connect; 2=Ack; 3=Message
	#[getter]
	fn get_msg_type(&self) -> PyResult<u8> {
		Ok(match self.inner {
			pkstl::IncomingBinaryMessage::Connect{..} => 1,
			pkstl::IncomingBinaryMessage::Ack{..} => 2,
			pkstl::IncomingBinaryMessage::Message{..} => 3,
		})
	}
	
	/// Message data
	/// If msg_type=Connect: [data, pubkey]
	/// Else: [data]
	#[getter]
	fn get_data(&self) -> PyResult<Vec<PyObject>> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		match &self.inner {
			pkstl::IncomingBinaryMessage::Connect{custom_datas, peer_sig_public_key} => match custom_datas {
				Some(val) => Ok(vec![
					PyBytes::new(py, val.as_slice()).to_object(py),
					PyBytes::new(py, peer_sig_public_key.as_slice()).to_object(py)
				]),
				None => Err(ValueError::py_err("Error"))
			},
			pkstl::IncomingBinaryMessage::Ack{custom_datas} => match custom_datas {
				Some(val) => Ok(vec![PyBytes::new(py, val.as_slice()).to_object(py)]),
				None => Err(ValueError::py_err("Error"))
			},
			pkstl::IncomingBinaryMessage::Message{datas} => match datas {
				Some(val) => Ok(vec![PyBytes::new(py, val.as_slice()).to_object(py)]),
				None => Err(ValueError::py_err("Error"))
			}
		}
	}
}

#[pyclass]
struct SecureLayer {
	inner: pkstl::SecureLayer
}

#[pymethods]
impl SecureLayer {
	
	/// Instanciate SecureLayer
	/// Arguments:
	/// - config: SecureLayerConfig
	/// - sig_key_pair_seed: Seed32 (optional)
	/// - expected_remote_sig_public_key: bytes (optional)
	/// Returns: SecureLayer
	#[staticmethod]
	fn create(
		config: &SecureLayerConfig,
		sig_key_pair_seed: Option<&Seed32>,
		expected_remote_sig_public_key: Option<Py<PyBytes>>
	) -> PyResult<SecureLayer>
	{
		let gil = Python::acquire_gil();
        let py = gil.python();
		match pkstl::SecureLayer::create(
			config.inner,
			match sig_key_pair_seed {
				Some(val) => Some(val.inner.clone()),
				None => None
			},
			match expected_remote_sig_public_key {
				Some(val) => Some(val.as_ref(py).as_bytes().to_vec()),
				None => None
			}
		) {
			Ok(secure_layer) => Ok(SecureLayer { inner: secure_layer }),
			Err(_) => Err(ValueError::py_err("Error"))
		}
	}
	
	/// Write message `Connect`
	/// Arguments:
	/// - custom_data: bytes (optional)
	/// Returns: bytes
	fn write_connect_msg_bin(&mut self, custom_data: Option<Py<PyBytes>>) -> PyResult<PyObject> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		let mut channel = BufWriter::new(Vec::with_capacity(1_024));
		
		match custom_data {
			Some(val) =>
				match self.inner.write_connect_msg_bin(
					Some(val.as_ref(py).as_bytes()),
					&mut channel
				) {
					Ok(_) => Ok(PyBytes::new(py, channel.buffer()).to_object(py)),
					Err(_) => Err(ValueError::py_err("Error"))
				},
			None =>
				match self.inner.write_connect_msg_bin(
					None,
					&mut channel
				) {
					Ok(_) => Ok(PyBytes::new(py, channel.buffer()).to_object(py)),
					Err(_) => Err(ValueError::py_err("Error"))
				}
		}
	}
	
	/// Write message `Ack`
	/// Arguments:
	/// - custom_data: bytes (optional)
	/// Returns: bytes
	fn write_ack_msg_bin(&mut self, custom_data: Option<Py<PyBytes>>) -> PyResult<PyObject> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		let mut channel = BufWriter::new(Vec::with_capacity(1_024));
		
		match custom_data {
			Some(val) =>
				match self.inner.write_ack_msg_bin(
					Some(val.as_ref(py).as_bytes()),
					&mut channel
				) {
					Ok(_) => Ok(PyBytes::new(py, channel.buffer()).to_object(py)),
					Err(_) => Err(ValueError::py_err("Error"))
				},
			None =>
				match self.inner.write_ack_msg_bin(
					None,
					&mut channel
				) {
					Ok(_) => Ok(PyBytes::new(py, channel.buffer()).to_object(py)),
					Err(_) => Err(ValueError::py_err("Error"))
				}
		}
	}
	
	/// Write
	/// Arguments:
	/// - data: bytes
	/// Returns: bytes
	fn write_bin(&mut self, data: Py<PyBytes>) -> PyResult<PyObject> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		let mut channel = BufWriter::new(Vec::with_capacity(1_024));
		
		match self.inner.write_bin(
			data.as_ref(py).as_bytes(),
			&mut channel
		) {
			Ok(_) => Ok(PyBytes::new(py, channel.buffer()).to_object(py)),
			Err(_) => Err(ValueError::py_err("Error"))
		}
	}
	
	/// Read binary messages
	/// Arguments:
	/// - incoming_data: bytes
	/// Returns: [IncomingBinaryMessage]
	fn read_bin(&mut self, incoming_data: Py<PyBytes>) -> PyResult<Vec<IncomingBinaryMessage>> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		let val = self.inner.read_bin(incoming_data.as_ref(py).as_bytes());
		match val {
			Ok(val) => Ok(val.into_iter().map(|x| IncomingBinaryMessage { inner: x }).collect()),
			Err(_) => Err(ValueError::py_err("Error"))
		}
	}
	
	/// Set config
	/// Arguments:
	/// - new_config: SecureLayerConfig
	fn change_config(&mut self, new_config: &SecureLayerConfig) -> PyResult<()> {
		match self.inner.change_config(new_config.inner) {
			Ok(_) => Ok(()),
			Err(_) => Err(ValueError::py_err("Error"))
		}
	}
}

#[pymodule]
fn pkstl(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
	m.add_class::<Ed25519KeyPair>()?;
	m.add_class::<Seed32>()?;
	m.add_class::<SecureLayerConfig>()?;
	m.add_class::<MinimalSecureLayer>()?;
	m.add_class::<IncomingBinaryMessage>()?;
	m.add_class::<SecureLayer>()?;
	
	Ok(())
}