From bea6e54a3710c629ea143cf1cdadca563adf66fd Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Sat, 21 Mar 2020 17:53:50 -0500
Subject: [PATCH] Make ssl work on windows

---
 Cargo.lock           | 11 +++++++++
 Lib/ssl.py           |  2 +-
 vm/Cargo.toml        |  3 ++-
 vm/src/obj/objset.rs | 11 +++++++++
 vm/src/stdlib/ssl.rs | 57 +++++++++++++++++++++++++++++++++++++++++++-
 5 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 227d2d1a9..b952e6628 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1639,6 +1639,7 @@ dependencies = [
  "rustpython-derive",
  "rustpython-parser",
  "rustyline",
+ "schannel",
  "serde",
  "serde_json",
  "sha-1",
@@ -1706,6 +1707,16 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bfa8506c1de11c9c4e4c38863ccbe02a305c8188e85a05a784c9e11e1c3910c8"
 
+[[package]]
+name = "schannel"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "039c25b130bd8c1321ee2d7de7fde2659fa9c2744e4bb29711cfc852ea53cd19"
+dependencies = [
+ "lazy_static 1.4.0",
+ "winapi",
+]
+
 [[package]]
 name = "semver"
 version = "0.9.0"
diff --git a/Lib/ssl.py b/Lib/ssl.py
index a578c0f55..0904a8a14 100644
--- a/Lib/ssl.py
+++ b/Lib/ssl.py
@@ -158,7 +158,7 @@ _SSLv2_IF_EXISTS = getattr(_SSLMethod, 'PROTOCOL_SSLv2', None)
 
 
 if sys.platform == "win32":
-    from _ssl import enum_certificates, enum_crls
+    from _ssl import enum_certificates #, enum_crls
 
 from socket import socket, AF_INET, SOCK_STREAM, create_connection
 from socket import SOL_SOCKET, SO_TYPE
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
index 0812c2400..4d1d0bdf6 100644
--- a/vm/Cargo.toml
+++ b/vm/Cargo.toml
@@ -101,10 +101,11 @@ libz-sys = "1.0"
 
 [target.'cfg(windows)'.dependencies]
 winreg = "0.7"
+schannel = "0.1"
 
 [target."cfg(windows)".dependencies.winapi]
 version = "0.3"
-features = ["winsock2", "handleapi", "ws2def", "std", "winbase"]
+features = ["winsock2", "handleapi", "ws2def", "std", "winbase", "wincrypt"]
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 wasm-bindgen = "0.2"
diff --git a/vm/src/obj/objset.rs b/vm/src/obj/objset.rs
index 369af668d..c9adfc944 100644
--- a/vm/src/obj/objset.rs
+++ b/vm/src/obj/objset.rs
@@ -631,6 +631,17 @@ macro_rules! multi_args_frozenset {
 
 #[pyimpl(flags(BASETYPE))]
 impl PyFrozenSet {
+    pub fn from_iter(
+        vm: &VirtualMachine,
+        it: impl IntoIterator<Item = PyObjectRef>,
+    ) -> PyResult<Self> {
+        let mut inner = PySetInner::default();
+        for elem in it {
+            inner.add(&elem, vm)?;
+        }
+        Ok(Self { inner })
+    }
+
     #[pyslot]
     fn tp_new(
         cls: PyClassRef,
diff --git a/vm/src/stdlib/ssl.rs b/vm/src/stdlib/ssl.rs
index 5ae84bc57..4b08f505d 100644
--- a/vm/src/stdlib/ssl.rs
+++ b/vm/src/stdlib/ssl.rs
@@ -106,6 +106,46 @@ fn obj2py(obj: &Asn1ObjectRef) -> PyNid {
     )
 }
 
+#[cfg(windows)]
+fn ssl_enum_certificates(store_name: PyStringRef, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
+    use crate::obj::objset::PyFrozenSet;
+    use schannel::{cert_context::ValidUses, cert_store::CertStore, RawPointer};
+    use winapi::um::wincrypt;
+    // TODO: check every store for it, not just 2 of them:
+    // https://github.com/python/cpython/blob/3.8/Modules/_ssl.c#L5603-L5610
+    let open_fns = [CertStore::open_current_user, CertStore::open_local_machine];
+    let stores = open_fns
+        .iter()
+        .filter_map(|open| open(store_name.as_str()).ok())
+        .collect::<Vec<_>>();
+    let certs = stores.iter().map(|s| s.certs()).flatten().map(|c| {
+        let cert = vm.ctx.new_bytes(c.to_der().to_owned());
+        let enc_type = unsafe {
+            let ptr = c.as_ptr() as wincrypt::PCCERT_CONTEXT;
+            (*ptr).dwCertEncodingType
+        };
+        let enc_type = match enc_type {
+            wincrypt::X509_ASN_ENCODING => vm.new_str("x509_asn".to_owned()),
+            wincrypt::PKCS_7_ASN_ENCODING => vm.new_str("pkcs_7_asn".to_owned()),
+            other => vm.new_int(other),
+        };
+        let usage = match c.valid_uses()? {
+            ValidUses::All => vm.new_bool(true),
+            ValidUses::Oids(oids) => {
+                PyFrozenSet::from_iter(vm, oids.into_iter().map(|oid| vm.new_str(oid)))
+                    .unwrap()
+                    .into_ref(vm)
+                    .into_object()
+            }
+        };
+        Ok(vm.ctx.new_tuple(vec![cert, enc_type, usage]))
+    });
+    let certs = certs
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|e| super::os::convert_io_error(vm, e))?;
+    Ok(vm.ctx.new_list(certs))
+}
+
 #[derive(FromArgs)]
 struct Txt2ObjArgs {
     #[pyarg(positional_or_keyword)]
@@ -518,7 +558,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         &vm.ctx.types.type_type,
         &vm.ctx.exceptions.os_error,
     );
-    py_module!(vm, "_ssl", {
+    let module = py_module!(vm, "_ssl", {
         "_SSLContext" => PySslContext::make_class(ctx),
         "_SSLSocket" => PySslSocket::make_class(ctx),
         "SSLError" => ssl_error,
@@ -570,5 +610,20 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "ALERT_DESCRIPTION_DECODE_ERROR" => ctx.new_int(sys::SSL_AD_DECODE_ERROR),
         "ALERT_DESCRIPTION_ILLEGAL_PARAMETER" => ctx.new_int(sys::SSL_AD_ILLEGAL_PARAMETER),
         "ALERT_DESCRIPTION_UNRECOGNIZED_NAME" => ctx.new_int(sys::SSL_AD_UNRECOGNIZED_NAME),
+    });
+
+    extend_module_platform_specific(&module, vm);
+
+    module
+}
+
+#[cfg(windows)]
+fn extend_module_platform_specific(module: &PyObjectRef, vm: &VirtualMachine) {
+    let ctx = &vm.ctx;
+    extend_module!(vm, module, {
+        "enum_certificates" => ctx.new_function(ssl_enum_certificates),
     })
 }
+
+#[cfg(not(windows))]
+fn extend_module_platform_specific(_module: &PyObjectRef, _vm: &VirtualMachine) {}