Misc. tweaks

astral-sh · Sep 28, 2023 · e2e94f7 · e2e94f7
1 parent 3b12507
commit e2e94f7
Show file tree

Hide file tree

Showing 3 changed files with 204 additions and 208 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_bandit/S505.py b/crates/ruff_linter/resources/test/fixtures/flake8_bandit/S505.py
@@ -7,11 +7,12 @@
 from Cryptodome.PublicKey import DSA as pycryptodomex_dsa
 from Cryptodome.PublicKey import RSA as pycryptodomex_rsa
 
-
-# Okay
+# OK
 dsa.generate_private_key(key_size=2048, backend=backends.default_backend())
 ec.generate_private_key(curve=ec.SECP384R1, backend=backends.default_backend())
-rsa.generate_private_key(public_exponent=65537, key_size=2048, backend=backends.default_backend())
+rsa.generate_private_key(
+    public_exponent=65537, key_size=2048, backend=backends.default_backend()
+)
 pycrypto_dsa.generate(bits=2048)
 pycrypto_rsa.generate(bits=2048)
 pycryptodomex_dsa.generate(bits=2048)
@@ -23,11 +24,13 @@
 pycrypto_rsa.generate(2048)
 pycryptodomex_dsa.generate(2048)
 pycryptodomex_rsa.generate(2048)
-#
-# # Errors
+
+# Errors
 dsa.generate_private_key(key_size=2047, backend=backends.default_backend())
 ec.generate_private_key(curve=ec.SECT163R2, backend=backends.default_backend())
-rsa.generate_private_key(public_exponent=65537, key_size=2047, backend=backends.default_backend())
+rsa.generate_private_key(
+    public_exponent=65537, key_size=2047, backend=backends.default_backend()
+)
 pycrypto_dsa.generate(bits=2047)
 pycrypto_rsa.generate(bits=2047)
 pycryptodomex_dsa.generate(bits=2047)
@@ -40,8 +43,12 @@
 pycryptodomex_dsa.generate(2047)
 pycryptodomex_rsa.generate(2047)
 
-# Don't crash when the size is variable
-rsa.generate_private_key(public_exponent=65537, key_size=some_key_size, backend=backends.default_backend())
+# Don't crash when the size is variable.
+rsa.generate_private_key(
+    public_exponent=65537, key_size=some_key_size, backend=backends.default_backend()
+)
 
-# Can't reliably know which curve was passed, in some cases like below
-ec.generate_private_key(curve=curves[self.curve]['create'](self.size), backend=backends.default_backend())
+# Can't reliably know which curve was passed, in some cases like below.
+ec.generate_private_key(
+    curve=curves[self.curve]["create"](self.size), backend=backends.default_backend()
+)
diff --git a/crates/ruff_linter/src/rules/flake8_bandit/rules/weak_cryptographic_key.rs b/crates/ruff_linter/src/rules/flake8_bandit/rules/weak_cryptographic_key.rs
@@ -1,58 +1,18 @@
-use once_cell::sync::Lazy;
+use std::fmt::{Display, Formatter};
+
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast, Constant, Expr, ExprAttribute, ExprCall};
 use ruff_text_size::{Ranged, TextRange};
-use rustc_hash::FxHashSet;
-use similar::DiffableStr;
-use std::fmt::{Display, Formatter};
 
 use crate::checkers::ast::Checker;
 
-static VULNERABLE_ELLIPTIC_CURVE_KEYS: Lazy<FxHashSet<&'static str>> =
-    Lazy::new(|| FxHashSet::from_iter(["SECP192R1", "SECT163K1", "SECT163R2"]));
-
-#[derive(Debug, PartialEq, Eq)]
-enum CryptographicKey {
-    Dsa { key_size: u16 },
-    Ec { algorithm: String },
-    Rsa { key_size: u16 },
-}
-
-impl CryptographicKey {
-    const fn minimum_key_size(&self) -> u16 {
-        match self {
-            Self::Dsa { .. } | Self::Rsa { .. } => 2048,
-            Self::Ec { .. } => 224,
-        }
-    }
-
-    fn is_vulnerable(&self) -> bool {
-        match self {
-            Self::Dsa { key_size } | Self::Rsa { key_size } => key_size < &self.minimum_key_size(),
-            Self::Ec { algorithm } => VULNERABLE_ELLIPTIC_CURVE_KEYS.contains(algorithm.as_str()),
-        }
-    }
-}
-
-impl Display for CryptographicKey {
-    fn fmt(&self, fmt: &mut Formatter) -> std::fmt::Result {
-        match self {
-            CryptographicKey::Dsa { .. } => fmt.write_str("DSA"),
-            CryptographicKey::Ec { .. } => fmt.write_str("EC"),
-            CryptographicKey::Rsa { .. } => fmt.write_str("RSA"),
-        }
-    }
-}
-
 /// ## What it does
-/// Checks for uses of `cryptographic key lengths known to be vulnerable.
+/// Checks for uses of cryptographic keys with vulnerable key sizes.
 ///
 /// ## Why is this bad?
-/// Small key lengths can easily be breakable, as computational power
-/// increases. For DSA and RSA keys, it is recommended to use key lengths equal
-/// or higher to 2048 bits. For EC, it is recommended to use curves higher than
-/// 224 bits.
+/// Small keys are easily breakable. For DSA and RSA, keys should be at least
+/// 2048 bits long. For EC, keys should be at least 224 bits long.
 ///
 /// ## Example
 /// ```python
@@ -88,77 +48,8 @@ impl Violation for WeakCryptographicKey {
     }
 }
 
-fn extract_int_argument(call: &ExprCall, name: &str, position: usize) -> Option<(u16, TextRange)> {
-    let Some(argument) = call.arguments.find_argument(name, position) else {
-        return None;
-    };
-    let Expr::Constant(ast::ExprConstant {
-        value: Constant::Int(i),
-        ..
-    }) = argument
-    else {
-        return None;
-    };
-    Some((i.as_u16()?, argument.range()))
-}
-
-fn extract_cryptographic_key(
-    checker: &mut Checker,
-    call: &ExprCall,
-) -> Option<(CryptographicKey, TextRange)> {
-    return checker
-        .semantic()
-        .resolve_call_path(&call.func)
-        .and_then(|call_path| match call_path.as_slice() {
-            ["cryptography", "hazmat", "primitives", "asymmetric", function, "generate_private_key"] => {
-                return match function.as_str() {
-                    Some("dsa") => {
-                        let Some((key_size, range)) = extract_int_argument(call, "key_size", 0) else {return None};
-                        return Some((CryptographicKey::Dsa { key_size }, range));
-                    },
-                    Some("rsa") => {
-                        let Some((key_size, range)) = extract_int_argument(call, "key_size", 1) else {return None};
-                        return Some((CryptographicKey::Rsa { key_size }, range));
-                    },
-                    Some("ec") => {
-                        let Some(argument) = call.arguments.find_argument("curve", 0) else { return None };
-                        let Expr::Attribute(ExprAttribute { attr, value, .. }) = argument else { return None };
-
-                        if checker
-                            .semantic()
-                            .resolve_call_path(value)
-                            .is_some_and(|call_path| matches!(call_path.as_slice(), ["cryptography", "hazmat", "primitives", "asymmetric", "ec"]))
-                        {
-                            return Some((CryptographicKey::Ec{algorithm: attr.as_str().to_string()}, argument.range()));
-                        }
-                        return None;
-                    },
-                    _ => None,
-                };
-            },
-            ["Crypto" | "Cryptodome", "PublicKey", function, "generate"] => {
-                return match function.as_str() {
-                    Some("DSA") => {
-                        let Some((key_size, range)) = extract_int_argument(call, "bits", 0) else {return None};
-                        return Some((CryptographicKey::Dsa { key_size }, range));
-                    },
-                    Some("RSA") => {
-                        let Some((key_size, range)) = extract_int_argument(call, "bits", 0) else {return None};
-                        return Some((CryptographicKey::Dsa { key_size }, range));
-                    },
-                    _ => None,
-                };
-            },
-            _ => None,
-        });
-}
-
 /// S505
 pub(crate) fn weak_cryptographic_key(checker: &mut Checker, call: &ExprCall) {
-    let Expr::Attribute(_) = call.func.as_ref() else {
-        return;
-    };
-
     let Some((cryptographic_key, range)) = extract_cryptographic_key(checker, call) else {
         return;
     };
@@ -170,3 +61,102 @@ pub(crate) fn weak_cryptographic_key(checker: &mut Checker, call: &ExprCall) {
         ));
     }
 }
+
+#[derive(Debug, PartialEq, Eq)]
+enum CryptographicKey {
+    Dsa { key_size: u16 },
+    Ec { algorithm: String },
+    Rsa { key_size: u16 },
+}
+
+impl CryptographicKey {
+    const fn minimum_key_size(&self) -> u16 {
+        match self {
+            Self::Dsa { .. } | Self::Rsa { .. } => 2048,
+            Self::Ec { .. } => 224,
+        }
+    }
+
+    fn is_vulnerable(&self) -> bool {
+        match self {
+            Self::Dsa { key_size } | Self::Rsa { key_size } => key_size < &self.minimum_key_size(),
+            Self::Ec { algorithm } => {
+                matches!(algorithm.as_str(), "SECP192R1" | "SECT163K1" | "SECT163R2")
+            }
+        }
+    }
+}
+
+impl Display for CryptographicKey {
+    fn fmt(&self, fmt: &mut Formatter) -> std::fmt::Result {
+        match self {
+            CryptographicKey::Dsa { .. } => fmt.write_str("DSA"),
+            CryptographicKey::Ec { .. } => fmt.write_str("EC"),
+            CryptographicKey::Rsa { .. } => fmt.write_str("RSA"),
+        }
+    }
+}
+
+fn extract_cryptographic_key(
+    checker: &mut Checker,
+    call: &ExprCall,
+) -> Option<(CryptographicKey, TextRange)> {
+    let call_path = checker.semantic().resolve_call_path(&call.func)?;
+    match call_path.as_slice() {
+        ["cryptography", "hazmat", "primitives", "asymmetric", function, "generate_private_key"] => {
+            match *function {
+                "dsa" => {
+                    let (key_size, range) = extract_int_argument(call, "key_size", 0)?;
+                    Some((CryptographicKey::Dsa { key_size }, range))
+                }
+                "rsa" => {
+                    let (key_size, range) = extract_int_argument(call, "key_size", 1)?;
+                    Some((CryptographicKey::Rsa { key_size }, range))
+                }
+                "ec" => {
+                    let argument = call.arguments.find_argument("curve", 0)?;
+                    let ExprAttribute { attr, value, .. } = argument.as_attribute_expr()?;
+                    let call_path = checker.semantic().resolve_call_path(value)?;
+                    if matches!(
+                        call_path.as_slice(),
+                        ["cryptography", "hazmat", "primitives", "asymmetric", "ec"]
+                    ) {
+                        Some((
+                            CryptographicKey::Ec {
+                                algorithm: attr.to_string(),
+                            },
+                            argument.range(),
+                        ))
+                    } else {
+                        None
+                    }
+                }
+                _ => None,
+            }
+        }
+        ["Crypto" | "Cryptodome", "PublicKey", function, "generate"] => match *function {
+            "DSA" => {
+                let (key_size, range) = extract_int_argument(call, "bits", 0)?;
+                Some((CryptographicKey::Dsa { key_size }, range))
+            }
+            "RSA" => {
+                let (key_size, range) = extract_int_argument(call, "bits", 0)?;
+                Some((CryptographicKey::Dsa { key_size }, range))
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+fn extract_int_argument(call: &ExprCall, name: &str, position: usize) -> Option<(u16, TextRange)> {
+    let argument = call.arguments.find_argument(name, position)?;
+    let Expr::Constant(ast::ExprConstant {
+        value: Constant::Int(i),
+        ..
+    }) = argument
+    else {
+        return None;
+    };
+    Some((i.as_u16()?, argument.range()))
+}