120 lines
3.6 KiB
Python
120 lines
3.6 KiB
Python
"""
|
|
Configuration management for DNSRecon tool.
|
|
Handles API key storage, rate limiting, and default settings.
|
|
"""
|
|
|
|
import os
|
|
from typing import Dict, Optional
|
|
|
|
|
|
class Config:
|
|
"""Configuration manager for DNSRecon application."""
|
|
|
|
def __init__(self):
|
|
"""Initialize configuration with default values."""
|
|
self.api_keys: Dict[str, Optional[str]] = {
|
|
'shodan': None,
|
|
'virustotal': None
|
|
}
|
|
|
|
# Default settings
|
|
self.default_recursion_depth = 2
|
|
self.default_timeout = 30
|
|
self.max_concurrent_requests = 5
|
|
self.large_entity_threshold = 100
|
|
|
|
# Rate limiting settings (requests per minute)
|
|
self.rate_limits = {
|
|
'crtsh': 60, # Free service, be respectful
|
|
'virustotal': 4, # Free tier limit
|
|
'shodan': 60, # API dependent
|
|
'dns': 100 # Local DNS queries
|
|
}
|
|
|
|
# Provider settings
|
|
self.enabled_providers = {
|
|
'crtsh': True, # Always enabled (free)
|
|
'dns': True, # Always enabled (free)
|
|
'virustotal': False, # Requires API key
|
|
'shodan': False # Requires API key
|
|
}
|
|
|
|
# Logging configuration
|
|
self.log_level = 'INFO'
|
|
self.log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
|
|
# Flask configuration
|
|
self.flask_host = '127.0.0.1'
|
|
self.flask_port = 5000
|
|
self.flask_debug = True
|
|
|
|
def set_api_key(self, provider: str, api_key: str) -> bool:
|
|
"""
|
|
Set API key for a provider.
|
|
|
|
Args:
|
|
provider: Provider name (shodan, virustotal)
|
|
api_key: API key string
|
|
|
|
Returns:
|
|
bool: True if key was set successfully
|
|
"""
|
|
if provider in self.api_keys:
|
|
self.api_keys[provider] = api_key
|
|
self.enabled_providers[provider] = True if api_key else False
|
|
return True
|
|
return False
|
|
|
|
def get_api_key(self, provider: str) -> Optional[str]:
|
|
"""
|
|
Get API key for a provider.
|
|
|
|
Args:
|
|
provider: Provider name
|
|
|
|
Returns:
|
|
API key or None if not set
|
|
"""
|
|
return self.api_keys.get(provider)
|
|
|
|
def is_provider_enabled(self, provider: str) -> bool:
|
|
"""
|
|
Check if a provider is enabled.
|
|
|
|
Args:
|
|
provider: Provider name
|
|
|
|
Returns:
|
|
bool: True if provider is enabled
|
|
"""
|
|
return self.enabled_providers.get(provider, False)
|
|
|
|
def get_rate_limit(self, provider: str) -> int:
|
|
"""
|
|
Get rate limit for a provider.
|
|
|
|
Args:
|
|
provider: Provider name
|
|
|
|
Returns:
|
|
Rate limit in requests per minute
|
|
"""
|
|
return self.rate_limits.get(provider, 60)
|
|
|
|
def load_from_env(self):
|
|
"""Load configuration from environment variables."""
|
|
if os.getenv('VIRUSTOTAL_API_KEY'):
|
|
self.set_api_key('virustotal', os.getenv('VIRUSTOTAL_API_KEY'))
|
|
|
|
if os.getenv('SHODAN_API_KEY'):
|
|
self.set_api_key('shodan', os.getenv('SHODAN_API_KEY'))
|
|
|
|
# Override default settings from environment
|
|
self.default_recursion_depth = int(os.getenv('DEFAULT_RECURSION_DEPTH', '2'))
|
|
self.flask_debug = os.getenv('FLASK_DEBUG', 'True').lower() == 'true'
|
|
self.default_timeout = 30
|
|
self.max_concurrent_requests = 5
|
|
|
|
|
|
# Global configuration instance
|
|
config = Config() |