From 4f37ff683dfaf66d33e46fe11c3c5d59549012e2 Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Wed, 17 Dec 2025 16:23:40 -0500 Subject: [PATCH 1/2] better caching --- Cargo.lock | 2 +- Cargo.toml | 2 +- README.md | 8 +- cloudcheck_update/__init__.py | 41 ------ pyproject.toml | 2 +- pytest.ini | 4 +- src/lib.rs | 260 ++++++++++++++++++++++------------ uv.lock | 2 +- 8 files changed, 178 insertions(+), 143 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1e9a5a3..7e708f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -56,7 +56,7 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cloudcheck" -version = "8.4.0" +version = "8.4.1" dependencies = [ "openssl", "pyo3", diff --git a/Cargo.toml b/Cargo.toml index 996ad89..36e4a60 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "cloudcheck" -version = "8.4.0" +version = "8.4.1" edition = "2024" description = "CloudCheck is a simple Rust tool to check whether an IP address or hostname belongs to a cloud provider." license = "GPL-3.0" diff --git a/README.md b/README.md index 8a17524..0a7d0fc 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,11 @@ [![PyPI](https://img.shields.io/pypi/v/cloudcheck)](https://pypi.org/project/cloudcheck/) [![Rust Version](https://img.shields.io/badge/rust-1.70+-orange)](https://www.rust-lang.org) [![Crates.io](https://img.shields.io/crates/v/cloudcheck?color=orange)](https://crates.io/crates/cloudcheck) -[![License](https://img.shields.io/badge/license-GPLv3-blue.svg)](https://github.com/blacklanternsecurity/cloudcheck/blob/master/LICENSE) +[![License](https://img.shields.io/badge/license-GPLv3-blue.svg)](https://github.com/blacklanternsecurity/cloudcheck/blob/stable/LICENSE) [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) -[![Rust Tests](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/rust-tests.yml/badge.svg?branch=master)](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/rust-tests.yml) -[![Python Tests](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/python-tests.yml/badge.svg?branch=master)](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/python-tests.yml) -[![Pipeline Tests](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/pipeline-tests.yml/badge.svg?branch=master)](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/pipeline-tests.yml) +[![Rust Tests](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/rust-tests.yml/badge.svg?branch=stable)](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/rust-tests.yml) +[![Python Tests](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/python-tests.yml/badge.svg?branch=stable)](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/python-tests.yml) +[![Pipeline Tests](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/pipeline-tests.yml/badge.svg?branch=stable)](https://github.com/blacklanternsecurity/cloudcheck/actions/workflows/pipeline-tests.yml) ### UPDATE: Now rewritten in Rust, with [34 supported cloud providers](#supported-cloud-providers)! diff --git a/cloudcheck_update/__init__.py b/cloudcheck_update/__init__.py index b0f70bd..29ceafa 100644 --- a/cloudcheck_update/__init__.py +++ b/cloudcheck_update/__init__.py @@ -11,47 +11,6 @@ logging.basicConfig(level=logging.INFO) log = logging.getLogger("cloudcheck") -# Dictionary to store loaded provider classes -_provider_classes: Dict[str, Type[BaseProvider]] = {} -# Dictionary to store instantiated providers -providers: Dict[str, BaseProvider] = {} - - -def load_provider_classes() -> Dict[str, Type[BaseProvider]]: - """Dynamically load all cloud provider classes from the providers directory.""" - global _provider_classes - - if _provider_classes: - return _provider_classes - - providers_path = Path(__file__).parent.parent / "cloudcheck" / "providers" - - for file in providers_path.glob("*.py"): - if file.stem in ("base", "__init__"): - continue - - try: - import_path = f"cloudcheck.providers.{file.stem}" - module = importlib.import_module(import_path) - - # Look for classes that inherit from BaseProvider - for attr_name in dir(module): - attr = getattr(module, attr_name) - if ( - isinstance(attr, type) - and issubclass(attr, BaseProvider) - and attr != BaseProvider - ): - provider_name = attr.__name__.lower() - _provider_classes[provider_name] = attr - print(f"Loaded provider class: {attr.__name__}") - - except Exception as e: - log.error(f"Failed to load provider from {file}: {e}") - raise - - return _provider_classes - project_root = Path(__file__).parent.parent json_path = project_root / "cloud_providers_v2.json" diff --git a/pyproject.toml b/pyproject.toml index 9720664..467e2a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "cloudcheck" -version = "8.4.0" +version = "8.4.1" description = "Detailed database of cloud providers. Instantly look up a domain or IP address" readme = "README.md" requires-python = ">=3.9" diff --git a/pytest.ini b/pytest.ini index 69d92a6..c81163f 100644 --- a/pytest.ini +++ b/pytest.ini @@ -3,5 +3,5 @@ testpaths = . python_files = test_*.py python_classes = Test* python_functions = test_* -norecursedirs = old -addopts = --ignore=old +norecursedirs = old target +addopts = --ignore=old --ignore=target diff --git a/src/lib.rs b/src/lib.rs index 1f38b67..67e2960 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,7 +4,7 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use std::time::{Duration, SystemTime}; -use tokio::sync::OnceCell; +use tokio::sync::{Mutex, RwLock}; #[cfg(feature = "py")] mod python; @@ -25,10 +25,14 @@ struct ProviderData { domains: Vec, } +type ProvidersMap = HashMap>; +type Error = Box; + #[derive(Clone)] pub struct CloudCheck { - radix: Arc>, - providers: Arc>>>, + radix: Arc>>, + providers: Arc>>, + last_fetch: Arc>>, } impl Default for CloudCheck { @@ -40,8 +44,9 @@ impl Default for CloudCheck { impl CloudCheck { pub fn new() -> Self { CloudCheck { - radix: Arc::new(OnceCell::new()), - providers: Arc::new(OnceCell::new()), + radix: Arc::new(RwLock::new(None)), + providers: Arc::new(RwLock::new(None)), + last_fetch: Arc::new(Mutex::new(None)), } } @@ -50,7 +55,7 @@ impl CloudCheck { .unwrap_or_else(|_| CLOUDCHECK_SIGNATURE_URL.to_string()) } - fn get_cache_path() -> Result> { + fn get_cache_path() -> Result { let home = std::env::var("HOME")?; let mut path = PathBuf::from(home); path.push(".cache"); @@ -61,7 +66,7 @@ impl CloudCheck { async fn fetch_and_cache( cache_path: &PathBuf, - ) -> Result> { + ) -> Result { let url = Self::get_signature_url(); let response = reqwest::get(&url).await?; let json_data = response.text().await?; @@ -74,96 +79,164 @@ impl CloudCheck { Ok(json_data) } - async fn ensure_loaded(&self) -> Result<(), Box> { - let radix_cell = Arc::clone(&self.radix); - let providers_cell = Arc::clone(&self.providers); - - radix_cell - .get_or_try_init(|| async { - let cache_path = Self::get_cache_path()?; - let cache_valid_duration = Duration::from_secs(24 * 60 * 60); - - let json_data = match tokio::fs::metadata(&cache_path).await { - Ok(metadata) => { - let modified = metadata.modified()?; - let now = SystemTime::now(); - - if let Ok(elapsed) = now.duration_since(modified) { - if elapsed < cache_valid_duration { - tokio::fs::read_to_string(&cache_path).await? - } else { - Self::fetch_and_cache(&cache_path).await? - } - } else { - Self::fetch_and_cache(&cache_path).await? - } - } - Err(_) => Self::fetch_and_cache(&cache_path).await?, - }; + /// Gets the last fetch time, checking in-memory timestamp first. + /// If no in-memory timestamp exists (first run), falls back to checking + /// the cache file's modification time. Returns None if file doesn't exist. + async fn get_last_fetch_time( + &self, + cache_path: &PathBuf, + ) -> Result, Error> { + let last_fetch = self.last_fetch.lock().await; + match *last_fetch { + Some(time) => Ok(Some(time)), + None => { + // No in-memory timestamp - check file modification time + drop(last_fetch); + match tokio::fs::metadata(cache_path).await { + Ok(metadata) => Ok(metadata.modified().ok()), + Err(_) => Ok(None), + } + } + } + } - let providers_data: HashMap = - serde_json::from_str(&json_data)?; - - let mut radix = RadixTarget::new(&[], ScopeMode::Normal)?; - let mut providers_map: HashMap> = HashMap::new(); - - for (_, provider) in providers_data { - let cloud_provider = CloudProvider { - name: provider.name.clone(), - tags: provider.tags.clone(), - }; - - for cidr in provider.cidrs { - let normalized = match radix.get(&cidr) { - Some(n) => n, - None => match radix.insert(&cidr) { - Ok(Some(n)) => n, - Ok(None) => continue, - Err(e) => { - eprintln!("Error inserting CIDR '{}': {}", cidr, e); - continue; - } - }, - }; - providers_map - .entry(normalized.clone()) - .or_default() - .push(cloud_provider.clone()); + /// Loads JSON data either from network (if refresh needed) or from cache file. + /// Returns (json_data, fetched_fresh) where fetched_fresh indicates if we + /// fetched from network. Sets last_fetch timestamp on first cache load to + /// track process runtime. Falls back to network fetch if cache read fails. + async fn load_json_data( + &self, + cache_path: &PathBuf, + needs_refresh: bool, + ) -> Result<(String, bool), Error> { + if needs_refresh { + let data = Self::fetch_and_cache(cache_path).await?; + Ok((data, true)) + } else { + match tokio::fs::read_to_string(cache_path).await { + Ok(data) => { + // First load from cache - set timestamp to track process runtime + let now = SystemTime::now(); + let mut last_fetch = self.last_fetch.lock().await; + if last_fetch.is_none() { + *last_fetch = Some(now); } + Ok((data, false)) + } + Err(_) => { + // Cache file was deleted between stat and read, fetch fresh + let data = Self::fetch_and_cache(cache_path).await?; + Ok((data, true)) + } + } + } + } - for domain in provider.domains { - // Clean domain: strip comments (everything after #) and trim whitespace - let cleaned_domain = domain.split('#').next().unwrap_or(&domain).trim(); - - if cleaned_domain.is_empty() { + /// Parses JSON and builds the radix tree and providers map. + /// For each provider, inserts all CIDRs and domains into the radix tree, + /// normalizing them in the process. Maps normalized values to provider lists. + fn build_data_structures( + json_data: &str, + ) -> Result<(RadixTarget, ProvidersMap), Error> { + let providers_data: HashMap = serde_json::from_str(json_data)?; + + let mut radix = RadixTarget::new(&[], ScopeMode::Normal)?; + let mut providers_map: ProvidersMap = HashMap::new(); + + for (_, provider) in providers_data { + let cloud_provider = CloudProvider { + name: provider.name.clone(), + tags: provider.tags.clone(), + }; + + // Insert all CIDRs for this provider + for cidr in provider.cidrs { + let normalized = match radix.get(&cidr) { + Some(n) => n, + None => match radix.insert(&cidr) { + Ok(Some(n)) => n, + Ok(None) => continue, + Err(e) => { + eprintln!("Error inserting CIDR '{}': {}", cidr, e); + continue; + } + }, + }; + providers_map + .entry(normalized.clone()) + .or_default() + .push(cloud_provider.clone()); + } + + // Insert all domains for this provider + for domain in provider.domains { + let normalized = match radix.get(&domain) { + Some(n) => n, + None => match radix.insert(&domain) { + Ok(Some(n)) => n, + Ok(None) => continue, + Err(e) => { + eprintln!("Error inserting domain '{}': {}", domain, e); continue; } + }, + }; + providers_map + .entry(normalized.clone()) + .or_default() + .push(cloud_provider.clone()); + } + } - let normalized = match radix.get(cleaned_domain) { - Some(n) => n, - None => match radix.insert(cleaned_domain) { - Ok(Some(n)) => n, - Ok(None) => continue, - Err(e) => { - eprintln!("Error inserting domain '{}': {}", cleaned_domain, e); - continue; - } - }, - }; - providers_map - .entry(normalized.clone()) - .or_default() - .push(cloud_provider.clone()); - } - } + Ok((radix, providers_map)) + } + + /// Ensures data is loaded and fresh. Checks if refresh is needed based on + /// 24-hour process runtime. Returns early if data is already loaded and fresh. + /// Otherwise loads data (from network or cache), builds structures, and updates + /// the in-memory timestamp if we fetched fresh data. + async fn ensure_loaded(&self) -> Result<(), Error> { + let cache_valid_duration = Duration::from_secs(24 * 60 * 60); + let now = SystemTime::now(); + let cache_path = Self::get_cache_path()?; + + // Check if we need refresh (uses in-memory timestamp, falls back to file stat) + let last_fetch_time = self.get_last_fetch_time(&cache_path).await?; + let needs_refresh = match last_fetch_time { + Some(fetch_time) => now + .duration_since(fetch_time) + .map(|elapsed| elapsed >= cache_valid_duration) + .unwrap_or(true), + None => true, + }; + + // Early return if data is already loaded and fresh + { + let radix_guard = self.radix.read().await; + if radix_guard.is_some() && !needs_refresh { + return Ok(()); + } + } + + // Load JSON data and build structures + let (json_data, fetched_fresh) = self.load_json_data(&cache_path, needs_refresh).await?; + let (radix, providers_map) = Self::build_data_structures(&json_data)?; - providers_cell - .set(providers_map) - .map_err(|_| "Failed to set providers")?; + // Update in-memory data structures + { + let mut radix_guard = self.radix.write().await; + *radix_guard = Some(radix); + } + { + let mut providers_guard = self.providers.write().await; + *providers_guard = Some(providers_map); + } - Ok::>(radix) - }) - .await?; + // Update timestamp if we fetched fresh data + if fetched_fresh { + let mut last_fetch = self.last_fetch.lock().await; + *last_fetch = Some(now); + } Ok(()) } @@ -171,11 +244,14 @@ impl CloudCheck { pub async fn lookup( &self, target: &str, - ) -> Result, Box> { + ) -> Result, Error> { self.ensure_loaded().await?; - let radix = self.radix.get().unwrap(); - let providers = self.providers.get().unwrap(); + let radix_guard = self.radix.read().await; + let providers_guard = self.providers.read().await; + + let radix = radix_guard.as_ref().unwrap(); + let providers = providers_guard.as_ref().unwrap(); if let Some(normalized) = radix.get(target) { Ok(providers.get(&normalized).cloned().unwrap_or_default()) diff --git a/uv.lock b/uv.lock index 92721f8..36d642f 100644 --- a/uv.lock +++ b/uv.lock @@ -140,7 +140,7 @@ wheels = [ [[package]] name = "cloudcheck" -version = "8.4.0" +version = "8.4.1" source = { editable = "." } [package.dev-dependencies] From 8975f6180856f33e8d3163d5fefa83ef8daa22e9 Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Wed, 17 Dec 2025 16:25:16 -0500 Subject: [PATCH 2/2] fmt --- src/lib.rs | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 67e2960..ebc83cf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -64,9 +64,7 @@ impl CloudCheck { Ok(path) } - async fn fetch_and_cache( - cache_path: &PathBuf, - ) -> Result { + async fn fetch_and_cache(cache_path: &PathBuf) -> Result { let url = Self::get_signature_url(); let response = reqwest::get(&url).await?; let json_data = response.text().await?; @@ -82,10 +80,7 @@ impl CloudCheck { /// Gets the last fetch time, checking in-memory timestamp first. /// If no in-memory timestamp exists (first run), falls back to checking /// the cache file's modification time. Returns None if file doesn't exist. - async fn get_last_fetch_time( - &self, - cache_path: &PathBuf, - ) -> Result, Error> { + async fn get_last_fetch_time(&self, cache_path: &PathBuf) -> Result, Error> { let last_fetch = self.last_fetch.lock().await; match *last_fetch { Some(time) => Ok(Some(time)), @@ -135,9 +130,7 @@ impl CloudCheck { /// Parses JSON and builds the radix tree and providers map. /// For each provider, inserts all CIDRs and domains into the radix tree, /// normalizing them in the process. Maps normalized values to provider lists. - fn build_data_structures( - json_data: &str, - ) -> Result<(RadixTarget, ProvidersMap), Error> { + fn build_data_structures(json_data: &str) -> Result<(RadixTarget, ProvidersMap), Error> { let providers_data: HashMap = serde_json::from_str(json_data)?; let mut radix = RadixTarget::new(&[], ScopeMode::Normal)?; @@ -241,10 +234,7 @@ impl CloudCheck { Ok(()) } - pub async fn lookup( - &self, - target: &str, - ) -> Result, Error> { + pub async fn lookup(&self, target: &str) -> Result, Error> { self.ensure_loaded().await?; let radix_guard = self.radix.read().await;