126 lines
3.8 KiB
Python
126 lines
3.8 KiB
Python
"""
|
|
Per-session configuration management for DNSRecon.
|
|
Provides isolated configuration instances for each user session.
|
|
"""
|
|
|
|
import os
|
|
from typing import Dict, Optional
|
|
|
|
|
|
class SessionConfig:
|
|
"""
|
|
Session-specific configuration that inherits from global config
|
|
but maintains isolated API keys and provider settings.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize session config with global defaults."""
|
|
# Copy all attributes from global config
|
|
self.api_keys: Dict[str, Optional[str]] = {
|
|
'shodan': None,
|
|
'virustotal': None
|
|
}
|
|
|
|
# Default settings (copied from global config)
|
|
self.default_recursion_depth = 2
|
|
self.default_timeout = 30
|
|
self.max_concurrent_requests = 5
|
|
self.large_entity_threshold = 100
|
|
|
|
# Rate limiting settings (per session)
|
|
self.rate_limits = {
|
|
'crtsh': 60,
|
|
'virustotal': 4,
|
|
'shodan': 60,
|
|
'dns': 100
|
|
}
|
|
|
|
# Provider settings (per session)
|
|
self.enabled_providers = {
|
|
'crtsh': True,
|
|
'dns': True,
|
|
'virustotal': False,
|
|
'shodan': False
|
|
}
|
|
|
|
# Logging configuration
|
|
self.log_level = 'INFO'
|
|
self.log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
|
|
# Flask configuration (shared)
|
|
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 in this session.
|
|
|
|
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 in this session.
|
|
|
|
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 in this session.
|
|
|
|
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 in this session.
|
|
|
|
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 (only if not already set)."""
|
|
if os.getenv('VIRUSTOTAL_API_KEY') and not self.api_keys['virustotal']:
|
|
self.set_api_key('virustotal', os.getenv('VIRUSTOTAL_API_KEY'))
|
|
|
|
if os.getenv('SHODAN_API_KEY') and not self.api_keys['shodan']:
|
|
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.default_timeout = 30
|
|
self.max_concurrent_requests = 5
|
|
|
|
|
|
def create_session_config() -> SessionConfig:
|
|
"""Create a new session configuration instance."""
|
|
session_config = SessionConfig()
|
|
session_config.load_from_env()
|
|
return session_config |