[Client]: add ca_cert_file parameter

deedy5 · Aug 14, 2024 · 156c5c6 · 156c5c6
1 parent f3290b6
commit 156c5c6
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 42 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -66,8 +66,8 @@ jobs:
           for version in 3.8 3.9 3.10 3.11 3.12; do
             uv venv --preview --python $version
             source .venv/bin/activate
+            uv pip install certifi pytest
             uv pip install primp --no-index --find-links dist --force-reinstall
-            uv pip install pytest
             pytest
           done
       - name: pytest
@@ -86,8 +86,8 @@ jobs:
             for version in 3.8 3.9 3.10 3.11 3.12; do
               uv venv --preview --python $version
               source .venv/bin/activate
+              uv pip install certifi pytest
               uv pip install primp --no-index --find-links dist --force-reinstall
-              uv pip install pytest
               pytest
             done
 
@@ -145,8 +145,8 @@ jobs:
             for version in 3.8 3.9 3.10 3.11 3.12; do
               python$version -m venv .venv
               source .venv/bin/activate
+              pip install certifi pytest
               pip install primp --no-index --find-links dist --force-reinstall
-              pip install pytest
               pytest
             done
 
@@ -187,8 +187,8 @@ jobs:
           for version in 3.8 3.9 3.10 3.11 3.12; do
             uv venv --preview --python $version
             source .venv/Scripts/activate
+            uv pip install certifi pytest
             uv pip install primp --no-index --find-links dist --force-reinstall
-            uv pip install pytest
             pytest
           done
 
@@ -226,8 +226,8 @@ jobs:
           for version in 3.8 3.9 3.10 3.11 3.12; do
             uv venv --preview --python $version
             source .venv/bin/activate
+            uv pip install certifi pytest
             uv pip install primp --no-index --find-links dist --force-reinstall
-            uv pip install pytest
             pytest
           done
 

diff --git a/README.md b/README.md
@@ -69,6 +69,7 @@ class Client:
         follow_redirects (bool, optional): Whether to follow redirects. Default is True.
         max_redirects (int, optional): Maximum redirects to follow. Default 20. Applies if `follow_redirects` is True.
         verify (bool, optional): Verify SSL certificates. Default is True.
+        ca_cert_file (str, optional): Path to CA certificate store. Default is None.
         http1 (bool, optional): Use only HTTP/1.1. Default is None.
         http2 (bool, optional): Use only HTTP/2. Default is None.
          
@@ -198,6 +199,12 @@ print(r.text)
 resp = primp.Client(proxy="http://127.0.0.1:8080").get("https://tls.peet.ws/api/all")
 print(resp.json())
 
+# Using custom CA certificate store: file or certifi.where()
+resp = primp.Client(ca_cert_file="/cert/cacert.pem").get("https://tls.peet.ws/api/all")
+print(resp.json())
+resp = primp.Client(ca_cert_file=certifi.where()).get("https://tls.peet.ws/api/all")
+print(resp.json())
+
 # You can also use convenience functions that use a default Client instance under the hood:
 # primp.get() | primp.head() | primp.options() | primp.delete() | primp.post() | primp.patch() | primp.put()
 # These functions can accept the `impersonate` parameter:

diff --git a/pyproject.toml b/pyproject.toml
@@ -31,6 +31,7 @@ dependencies = []
 
 [project.optional-dependencies]
 dev = [
+    "certifi",
     "pytest>=8.1.1",
 ]
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -65,6 +65,7 @@ impl Client {
     /// * `follow_redirects` - A boolean to enable or disable following redirects. Default is `true`.
     /// * `max_redirects` - The maximum number of redirects to follow. Default is 20. Applies if `follow_redirects` is `true`.
     /// * `verify` - An optional boolean indicating whether to verify SSL certificates. Default is `true`.
+    /// * `ca_cert_file` - Path to CA certificate store. Default is None.
     /// * `http1` - An optional boolean indicating whether to use only HTTP/1.1. Default is `false`.
     /// * `http2` - An optional boolean indicating whether to use only HTTP/2. Default is `false`.
     ///
@@ -85,14 +86,16 @@ impl Client {
     ///     impersonate="chrome_123",
     ///     follow_redirects=True,
     ///     max_redirects=1,
-    ///     verify=False,
+    ///     verify=True,
+    ///     ca_cert_file="/cert/cacert.pem",
     ///     http1=True,
     ///     http2=False,
     /// )
     /// ```
     #[new]
-    #[pyo3(signature = (auth=None, auth_bearer=None, params=None, headers=None, cookies=None, cookie_store=None, referer=None, 
-        proxy=None, timeout=None, impersonate=None, follow_redirects=None, max_redirects=None, verify=None, http1=None, http2=None))]
+    #[pyo3(signature = (auth=None, auth_bearer=None, params=None, headers=None, cookies=None, 
+        cookie_store=None, referer=None, proxy=None, timeout=None, impersonate=None, follow_redirects=None, 
+        max_redirects=None, verify=None, ca_cert_file=None, http1=None, http2=None))]
     fn new(
         auth: Option<(String, Option<String>)>,
         auth_bearer: Option<String>,
@@ -107,6 +110,7 @@ impl Client {
         follow_redirects: Option<bool>,
         max_redirects: Option<usize>,
         verify: Option<bool>,
+        ca_cert_file: Option<String>,
         http1: Option<bool>,
         http2: Option<bool>,
     ) -> Result<Self> {
@@ -172,6 +176,11 @@ impl Client {
             client_builder = client_builder.danger_accept_invalid_certs(true);
         }
 
+        // Ca_cert_file
+        if let Some(ca_cert_file) = ca_cert_file {
+            client_builder = client_builder.ca_cert_file(ca_cert_file);
+        }
+
         // Http version: http1 || http2
         match (http1, http2) {
             (Some(true), Some(true)) => return Err(anyhow!("Both http1 and http2 cannot be true")),
@@ -601,7 +610,8 @@ impl Client {
 /// Convenience functions that use a default Client instance under the hood
 #[pyfunction]
 #[pyo3(signature = (method, url, params=None, headers=None, cookies=None, content=None, data=None, 
-    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None))]
+    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None,
+    ca_cert_file=None))]
 fn request(
     py: Python,
     method: &str,
@@ -618,6 +628,7 @@ fn request(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -633,6 +644,7 @@ fn request(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;
@@ -655,7 +667,7 @@ fn request(
 
 #[pyfunction]
 #[pyo3(signature = (url, params=None, headers=None, cookies=None, auth=None, auth_bearer=None, 
-    timeout=None, impersonate=None, verify=None))]
+    timeout=None, impersonate=None, verify=None, ca_cert_file=None))]
 fn get(
     py: Python,
     url: &str,
@@ -667,6 +679,7 @@ fn get(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -682,6 +695,7 @@ fn get(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;
@@ -699,7 +713,7 @@ fn get(
 
 #[pyfunction]
 #[pyo3(signature = (url, params=None, headers=None, cookies=None, auth=None, auth_bearer=None, 
-    timeout=None, impersonate=None, verify=None))]
+    timeout=None, impersonate=None, verify=None, ca_cert_file=None))]
 fn head(
     py: Python,
     url: &str,
@@ -711,6 +725,7 @@ fn head(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -726,6 +741,7 @@ fn head(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;
@@ -743,7 +759,7 @@ fn head(
 
 #[pyfunction]
 #[pyo3(signature = (url, params=None, headers=None, cookies=None, auth=None, auth_bearer=None, 
-    timeout=None, impersonate=None, verify=None))]
+    timeout=None, impersonate=None, verify=None, ca_cert_file=None))]
 fn options(
     py: Python,
     url: &str,
@@ -755,6 +771,7 @@ fn options(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -770,6 +787,7 @@ fn options(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;
@@ -787,7 +805,7 @@ fn options(
 
 #[pyfunction]
 #[pyo3(signature = (url, params=None, headers=None, cookies=None, auth=None, auth_bearer=None, 
-    timeout=None, impersonate=None, verify=None))]
+    timeout=None, impersonate=None, verify=None, ca_cert_file=None))]
 fn delete(
     py: Python,
     url: &str,
@@ -799,6 +817,7 @@ fn delete(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -814,6 +833,7 @@ fn delete(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;
@@ -831,7 +851,8 @@ fn delete(
 
 #[pyfunction]
 #[pyo3(signature = (url, params=None, headers=None, cookies=None, content=None, data=None, 
-    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None))]
+    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None,
+    ca_cert_file=None))]
 fn post(
     py: Python,
     url: &str,
@@ -847,6 +868,7 @@ fn post(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -862,6 +884,7 @@ fn post(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;
@@ -883,7 +906,8 @@ fn post(
 
 #[pyfunction]
 #[pyo3(signature = (url, params=None, headers=None, cookies=None, content=None, data=None, 
-    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None))]
+    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None,
+    ca_cert_file=None))]
 fn put(
     py: Python,
     url: &str,
@@ -899,6 +923,7 @@ fn put(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -914,6 +939,7 @@ fn put(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;
@@ -935,7 +961,8 @@ fn put(
 
 #[pyfunction]
 #[pyo3(signature = (url, params=None, headers=None, cookies=None, content=None, data=None, 
-    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None))]
+    json=None, files=None, auth=None, auth_bearer=None, timeout=None, impersonate=None, verify=None,
+    ca_cert_file=None))]
 fn patch(
     py: Python,
     url: &str,
@@ -951,6 +978,7 @@ fn patch(
     timeout: Option<f64>,
     impersonate: Option<&str>,
     verify: Option<bool>,
+    ca_cert_file: Option<String>,
 ) -> Result<Response> {
     let client = Client::new(
         None,
@@ -966,6 +994,7 @@ fn patch(
         None,
         None,
         verify,
+        ca_cert_file,
         None,
         None,
     )?;