Merge pull request #54 from Intreecom/feature/tls

README.md

@@ -480,4 +480,5 @@ To make built `Select` query return paginated iterator, add paged parameter in e
     rows = await Select("test").execute(scylla, paged=True)
     async for row in rows:
         print(row['id'])
-```
+```
+

python/scyllapy/__init__.py

@@ -1,5 +1,6 @@
 from importlib.metadata import version
 
+from . import extra_types
 from ._internal import (
     Batch,
     BatchType,
@@ -11,6 +12,7 @@
     QueryResult,
     Scylla,
     SerialConsistency,
+    SSLVerifyMode,
 )
 
 __version__ = version("scyllapy")
@@ -25,6 +27,7 @@
     "Batch",
     "BatchType",
     "QueryResult",
+    "SSLVerifyMode",
     "extra_types",
     "InlineBatch",
     "ExecutionProfile",

python/scyllapy/_internal/__init__.pyi

@@ -14,6 +14,18 @@ from scyllapy._internal.load_balancing import LoadBalancingPolicy
 _T = TypeVar("_T")
 _T2 = TypeVar("_T2")
 
+class SSLVerifyMode:
+    """
+    SSL Verify modes.
+
+    Used for SSL/TLS connection verification.
+    Read mode: https://docs.rs/openssl/0.10.68/openssl/ssl/struct.SslVerifyMode.html
+    """
+
+    NONE: SSLVerifyMode
+    PEER: SSLVerifyMode
+    FAIL_IF_NO_PEER_CERT: SSLVerifyMode
+
 class Scylla:
     """
     Scylla class.
@@ -32,6 +44,9 @@ class Scylla:
         password: str | None = None,
         keyspace: str | None = None,
         ssl_cert: str | None = None,
+        ssl_key: str | None = None,
+        ssl_ca_file: str | None = None,
+        ssl_verify_mode: SSLVerifyMode | None = None,
         conn_timeout: int | None = None,
         write_coalescing: bool | None = None,
         pool_size_per_host: int | None = None,
@@ -50,8 +65,13 @@ class Scylla:
             ["192.168.1.1:9042", "my_keyspace.node:9042"]
         :param username: Plain text auth username.
         :param password: Plain text auth password.
-        :param ssl_cert: Certficiate string to use
-            for connection. AWS requires it.
+        :param ssl_cert: Certficiate string to use for connection.
+            Should be PEM encoded x509 certificate string.
+        :param ssl_key: Key string to use for connection.
+            Should be RSA private key PEM encoded string.
+        :param ssl_ca_file: CA file to use for connection. This parameter
+            should be a path to the CA file (which is PEM encoded CA).
+        :param ssl_verify_mode: tells server on how to validate client's certificate.
         :param conn_timeout: Timeout in seconds.
         :param write_coalescing:
             If true, the driver will inject a small delay before flushing data

src/lib.rs

@@ -21,6 +21,7 @@ use crate::utils::add_submodule;
 fn _internal(py: Python<'_>, pymod: &PyModule) -> PyResult<()> {
     pyo3_log::init();
     pymod.add_class::<scylla_cls::Scylla>()?;
+    pymod.add_class::<scylla_cls::ScyllaPySSLVerifyMode>()?;
     pymod.add_class::<consistencies::ScyllaPyConsistency>()?;
     pymod.add_class::<consistencies::ScyllaPySerialConsistency>()?;
     pymod.add_class::<queries::ScyllaPyQuery>()?;

src/scylla_cls.rs

@@ -9,12 +9,37 @@ use crate::{
     utils::{parse_python_query_params, scyllapy_future},
 };
 use openssl::{
+    pkey::PKey,
+    rsa::Rsa,
     ssl::{SslContextBuilder, SslMethod, SslVerifyMode},
     x509::X509,
 };
 use pyo3::{pyclass, pymethods, PyAny, Python};
 use scylla::{frame::value::ValueList, prepared_statement::PreparedStatement, query::Query};
 
+/// SSL verification mode.
+#[pyclass(name = "SSLVerifyMode")]
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+#[allow(non_camel_case_types)]
+pub enum ScyllaPySSLVerifyMode {
+    /// No verification.
+    NONE,
+    // Verify peer.
+    PEER,
+    // Fail if no peer certificate.
+    FAIL_IF_NO_PEER_CERT,
+}
+
+impl From<ScyllaPySSLVerifyMode> for SslVerifyMode {
+    fn from(value: ScyllaPySSLVerifyMode) -> Self {
+        match value {
+            ScyllaPySSLVerifyMode::NONE => Self::NONE,
+            ScyllaPySSLVerifyMode::PEER => Self::PEER,
+            ScyllaPySSLVerifyMode::FAIL_IF_NO_PEER_CERT => Self::FAIL_IF_NO_PEER_CERT,
+        }
+    }
+}
+
 #[pyclass(frozen, weakref)]
 #[derive(Clone)]
 pub struct Scylla {
@@ -23,6 +48,9 @@ pub struct Scylla {
     password: Option<String>,
     keyspace: Option<String>,
     ssl_cert: Option<String>,
+    ssl_key: Option<String>,
+    ssl_ca_file: Option<String>,
+    ssl_verify_mode: Option<ScyllaPySSLVerifyMode>,
     connection_timeout: Option<u64>,
     write_coalescing: Option<bool>,
     disallow_shard_aware_port: Option<bool>,
@@ -110,6 +138,9 @@ impl Scylla {
         password = None,
         keyspace = None,
         ssl_cert = None,
+        ssl_key  = None,
+        ssl_ca_file = None,
+        ssl_verify_mode = None,
         connection_timeout = None,
         write_coalescing = None,
         pool_size_per_host = None,
@@ -128,6 +159,9 @@ impl Scylla {
         password: Option<String>,
         keyspace: Option<String>,
         ssl_cert: Option<String>,
+        ssl_key: Option<String>,
+        ssl_ca_file: Option<String>,
+        ssl_verify_mode: Option<ScyllaPySSLVerifyMode>,
         connection_timeout: Option<u64>,
         write_coalescing: Option<bool>,
         pool_size_per_host: Option<NonZeroUsize>,
@@ -144,6 +178,9 @@ impl Scylla {
             username,
             password,
             ssl_cert,
+            ssl_key,
+            ssl_ca_file,
+            ssl_verify_mode,
             keyspace,
             connection_timeout,
             write_coalescing,
@@ -177,8 +214,24 @@ impl Scylla {
         if let Some(cert_data) = self.ssl_cert.clone() {
             let mut ssl_context_builder = SslContextBuilder::new(SslMethod::tls())?;
             let pem = X509::from_pem(cert_data.as_bytes())?;
+            // If we have private key, we need to add it to context.
+            if let Some(key) = &self.ssl_key {
+                let rsa = Rsa::private_key_from_pem(key.as_bytes())?;
+                let pkey = PKey::from_rsa(rsa)?;
+                ssl_context_builder.set_private_key(&pkey)?;
+            }
+            // If we have CA file, we need to add it to context.
+            if let Some(ca_file) = &self.ssl_ca_file {
+                ssl_context_builder.set_ca_file(ca_file)?;
+            }
+            // Set verification mode.
+            // If no mode was passed, we use NONE.
+            let verify_mode = self
+                .ssl_verify_mode
+                .map_or(SslVerifyMode::NONE, SslVerifyMode::from);
+
             ssl_context_builder.set_certificate(&pem)?;
-            ssl_context_builder.set_verify(SslVerifyMode::NONE);
+            ssl_context_builder.set_verify(verify_mode);
             ssl_context = Some(ssl_context_builder.build());
         }
         let keyspace = self.keyspace.clone();