znippy_plugin_python/
lib.rs

1pub mod wheel;
2
3#[cfg(feature = "host-decompressors")]
4mod native;
5#[cfg(feature = "host-decompressors")]
6pub use native::NativePythonPlugin;
7
8use wheel::{WheelInfo, parse_wheel_filename};
9
10/// Extract Python package metadata from a wheel filename.
11/// No need to open the file — all info is encoded in the name per PEP 427.
12pub fn extract_python_metadata(path: &str, _data: &[u8]) -> Option<WheelInfo> {
13    let filename = path.rsplit('/').next().unwrap_or(path);
14    parse_wheel_filename(filename)
15}
16
17/// Extract METADATA content from inside a wheel (zip) for PEP 658 dependency resolution.
18/// Returns the raw bytes of the METADATA file if found.
19pub fn extract_metadata_from_wheel(data: &[u8]) -> Option<Vec<u8>> {
20    // METADATA lives at: {name}-{version}.dist-info/METADATA
21    // Use substring filter since we don't know the exact dist-info prefix
22    find_file_in_zip(data, "METADATA")
23}
24
25/// Extract RECORD from inside a wheel for file listing.
26pub fn extract_record_from_wheel(data: &[u8]) -> Option<Vec<u8>> {
27    find_file_in_zip(data, "RECORD")
28}
29
30fn find_file_in_zip(data: &[u8], needle: &str) -> Option<Vec<u8>> {
31    // Minimal ZIP central directory scan — no lzip dependency needed for fallback
32    let eocd = find_eocd(data)?;
33    let entry_count = u16::from_le_bytes(data[eocd + 8..eocd + 10].try_into().ok()?) as usize;
34    let cd_offset = u32::from_le_bytes(data[eocd + 16..eocd + 20].try_into().ok()?) as usize;
35
36    let mut pos = cd_offset;
37    for _ in 0..entry_count {
38        if pos + 46 > data.len() {
39            break;
40        }
41        if &data[pos..pos + 4] != b"PK\x01\x02" {
42            break;
43        }
44        let method = u16::from_le_bytes(data[pos + 10..pos + 12].try_into().ok()?);
45        let comp_size = u32::from_le_bytes(data[pos + 20..pos + 24].try_into().ok()?) as usize;
46        let uncomp_size = u32::from_le_bytes(data[pos + 24..pos + 28].try_into().ok()?) as usize;
47        let name_len = u16::from_le_bytes(data[pos + 28..pos + 30].try_into().ok()?) as usize;
48        let extra_len = u16::from_le_bytes(data[pos + 30..pos + 32].try_into().ok()?) as usize;
49        let comment_len = u16::from_le_bytes(data[pos + 32..pos + 34].try_into().ok()?) as usize;
50        let local_off = u32::from_le_bytes(data[pos + 42..pos + 46].try_into().ok()?) as usize;
51
52        let name = data
53            .get(pos + 46..pos + 46 + name_len)
54            .and_then(|b| std::str::from_utf8(b).ok())
55            .unwrap_or("");
56
57        // Match: ends with /METADATA or /RECORD (inside .dist-info/)
58        if name.ends_with(&format!("/{}", needle)) && name.contains(".dist-info/") {
59            if local_off + 30 > data.len() {
60                return None;
61            }
62            if &data[local_off..local_off + 4] != b"PK\x03\x04" {
63                return None;
64            }
65            let lname_len =
66                u16::from_le_bytes(data[local_off + 26..local_off + 28].try_into().ok()?)
67                    as usize;
68            let lextra_len =
69                u16::from_le_bytes(data[local_off + 28..local_off + 30].try_into().ok()?)
70                    as usize;
71            let data_start = local_off + 30 + lname_len + lextra_len;
72            let raw = data.get(data_start..data_start + comp_size)?;
73            return match method {
74                0 => Some(raw.to_vec()),
75                8 => {
76                    let mut out = vec![0u8; uncomp_size];
77                    let result = miniz_oxide::inflate::decompress_to_vec(raw).ok()?;
78                    out = result;
79                    Some(out)
80                }
81                _ => None,
82            };
83        }
84
85        pos += 46 + name_len + extra_len + comment_len;
86    }
87    None
88}
89
90fn find_eocd(data: &[u8]) -> Option<usize> {
91    const MIN_EOCD: usize = 22;
92    if data.len() < MIN_EOCD {
93        return None;
94    }
95    let earliest = data.len().saturating_sub(MIN_EOCD + 65535);
96    for i in (earliest..=data.len() - MIN_EOCD).rev() {
97        if data[i..].starts_with(b"PK\x05\x06") {
98            return Some(i);
99        }
100    }
101    None
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn test_extract_metadata_from_path() {
110        let info = extract_python_metadata(
111            "packages/numpy/numpy-1.26.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
112            &[],
113        ).unwrap();
114        assert_eq!(info.name, "numpy");
115        assert_eq!(info.version, "1.26.0");
116        assert_eq!(info.python_tag, "cp311");
117        assert_eq!(info.abi_tag, "cp311");
118        assert!(info.platform_tag.contains("manylinux"));
119    }
120
121    #[test]
122    fn test_sdist_not_wheel() {
123        let info = extract_python_metadata("packages/requests/requests-2.31.0.tar.gz", &[]);
124        assert!(info.is_none());
125    }
126}
znippy_plugin_python/lib.rs

znippy_plugin_python/
lib.rs
