src/usfx/modules/base.py (4739B)
1 """
2 Base Module Class
3
4 Abstract base class for all subdomain enumeration modules.
5 Provides common functionality and interface contract.
6 """
7
8 import logging
9 import threading
10 from abc import ABC, abstractmethod
11 from typing import Callable, Dict, Optional, TYPE_CHECKING
12
13 if TYPE_CHECKING:
14 from ..utils.dns_resolver import InternalDNSResolver
15
16 logger = logging.getLogger(__name__)
17
18
19 class BaseModule(ABC):
20 """
21 Abstract base class for subdomain enumeration modules.
22
23 All modules must implement the enumerate() method and provide
24 a MODULE_NAME class attribute.
25
26 Attributes:
27 MODULE_NAME: Unique identifier for the module
28 domain: Target domain
29 resolver: DNS resolver for lookups
30 discovered: Dict of discovered {subdomain: ip}
31 """
32
33 MODULE_NAME: str = "base"
34
35 def __init__(
36 self,
37 domain: str,
38 resolver: Optional['InternalDNSResolver'] = None,
39 timeout: float = 3.0,
40 cancel_check: Optional[Callable[[], bool]] = None
41 ):
42 """
43 Initialize the module.
44
45 Args:
46 domain: Target domain (e.g., 'corp.local')
47 resolver: Optional custom DNS resolver
48 timeout: Default timeout for operations
49 cancel_check: Optional function that returns True if cancelled
50 """
51 self.domain = domain.lower().strip('.')
52 self.timeout = timeout
53 self.discovered: Dict[str, Optional[str]] = {}
54 self._cancel_check = cancel_check
55 self._cancelled = False
56
57 # Use provided resolver or create a new one
58 if resolver:
59 self.resolver = resolver
60 else:
61 from ..utils.dns_resolver import InternalDNSResolver
62 self.resolver = InternalDNSResolver(timeout=timeout)
63
64 def is_cancelled(self) -> bool:
65 """Check if the module should stop"""
66 if self._cancelled:
67 return True
68 if self._cancel_check and self._cancel_check():
69 self._cancelled = True
70 return True
71 return False
72
73 def cancel(self):
74 """Cancel the module"""
75 self._cancelled = True
76
77 @abstractmethod
78 def enumerate(self, **kwargs) -> Dict[str, Optional[str]]:
79 """
80 Perform subdomain enumeration.
81
82 This method must be implemented by all modules.
83
84 Returns:
85 Dict of {subdomain: ip_address}
86 IP address may be None if not resolved.
87 """
88 pass
89
90 def _is_valid_subdomain(self, subdomain: str) -> bool:
91 """
92 Check if a discovered subdomain is valid for the target domain.
93
94 Args:
95 subdomain: Subdomain to validate
96
97 Returns:
98 True if valid subdomain of target domain
99 """
100 subdomain = subdomain.lower().rstrip('.')
101
102 # Must end with .domain
103 if not subdomain.endswith(f".{self.domain}"):
104 return False
105
106 # Must not be exactly the domain
107 if subdomain == self.domain:
108 return False
109
110 # Check for empty labels
111 labels = subdomain.split('.')
112 for label in labels:
113 if not label:
114 return False
115
116 return True
117
118 def _add_discovered(self, subdomain: str, ip: Optional[str] = None) -> bool:
119 """
120 Add a discovered subdomain to results.
121
122 Args:
123 subdomain: Discovered subdomain
124 ip: Optional IP address
125
126 Returns:
127 True if added (new), False if already existed
128 """
129 subdomain = subdomain.lower().rstrip('.')
130
131 if not self._is_valid_subdomain(subdomain):
132 return False
133
134 if subdomain not in self.discovered:
135 self.discovered[subdomain] = ip
136 logger.debug(f"[{self.MODULE_NAME}] Found: {subdomain} -> {ip}")
137 return True
138
139 # Update IP if we have one now
140 if ip and self.discovered[subdomain] is None:
141 self.discovered[subdomain] = ip
142
143 return False
144
145 def _resolve_subdomain(self, subdomain: str) -> Optional[str]:
146 """
147 Resolve a subdomain using the configured resolver.
148
149 Args:
150 subdomain: Subdomain to resolve
151
152 Returns:
153 IP address or None
154 """
155 return self.resolver.resolve_a(subdomain)
156
157 def get_results(self) -> Dict[str, Optional[str]]:
158 """
159 Get the discovered subdomains.
160
161 Returns:
162 Dict of {subdomain: ip_address}
163 """
164 return self.discovered.copy()
165
166 def get_count(self) -> int:
167 """
168 Get the number of discovered subdomains.
169
170 Returns:
171 Count of discovered subdomains
172 """
173 return len(self.discovered)