mirror of
https://github.com/yogeshojha/rengine.git
synced 2026-10-21 07:41:28 +02:00
Yogesh Ojha
69239333bc
- Added `clean_name` function to `validation.py` for name validation with max length. - Simplified docstrings in `dnsx` client, parser, and service files for clarity. - Enhanced error handling in `ViewDNSService` to prevent disabling keys on rate limits. - Improved WHOIS service to handle concurrent inserts gracefully with integrity checks. - Updated `Settings` class in worker config to inherit from `BaseAppSettings`. - Added new dependencies in `pyproject.toml` for `apprise` and `cryptography`.
295 lines
9.8 KiB
Python
295 lines
9.8 KiB
Python
"""DNS lookup service - standalone/persistent lookup plus cross-target correlation queries."""
|
|
|
|
import uuid
|
|
from collections import Counter
|
|
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from shared.enums.dns import DnsRecordType
|
|
from shared.logging import get_logger
|
|
from shared.models.dns import DnsLookup, DnsLookupSummary, DnsRecord
|
|
from shared.models.target import Target
|
|
from shared.utils.datetime import utc_now
|
|
from shared.utils.infra import is_shared_nameserver
|
|
from tools.dnsx.client import DnsxClient, DnsxError
|
|
from tools.dnsx.models import DnsxReconResponse
|
|
from tools.dnsx.parser import parse_dnsx_jsonl
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
# cap rows per correlation query
|
|
MAX_CORRELATION_TARGETS = 500
|
|
|
|
|
|
class DnsxServiceError(Exception):
|
|
"""Base exception for DNS service errors."""
|
|
|
|
|
|
class DnsxLookupError(DnsxServiceError):
|
|
"""Raised when a DNS lookup fails."""
|
|
|
|
|
|
class DnsxService:
|
|
"""DNS lookup service with optional DB persistence."""
|
|
|
|
def __init__(
|
|
self,
|
|
timeout: int = 120,
|
|
retry: int = 3,
|
|
threads: int = 10,
|
|
resolvers: list[str] | None = None,
|
|
) -> None:
|
|
self._timeout = timeout
|
|
self._retry = retry
|
|
self._threads = threads
|
|
self._resolvers = resolvers
|
|
self._client: DnsxClient | None = None
|
|
|
|
@property
|
|
def client(self) -> DnsxClient:
|
|
"""Lazy-init the dnsx client."""
|
|
if self._client is None:
|
|
try:
|
|
self._client = DnsxClient(
|
|
timeout=self._timeout,
|
|
retry=self._retry,
|
|
threads=self._threads,
|
|
resolvers=self._resolvers,
|
|
)
|
|
except DnsxError as e:
|
|
raise DnsxServiceError(str(e)) from e
|
|
return self._client
|
|
|
|
def do_recon(self, domain: str) -> DnsxReconResponse:
|
|
"""Run full DNS recon on a single domain (no DB)."""
|
|
result = self.client.recon(domain)
|
|
|
|
if not result.success and not result.has_output:
|
|
msg = f"dnsx recon failed for {domain}: {result.error or 'no output'}"
|
|
raise DnsxLookupError(msg)
|
|
|
|
parsed = parse_dnsx_jsonl(result.json_records)
|
|
if not parsed:
|
|
msg = f"dnsx returned no parseable results for {domain}"
|
|
raise DnsxLookupError(msg)
|
|
|
|
return parsed[0]
|
|
|
|
def do_recon_batch(self, domains: list[str]) -> list[DnsxReconResponse]:
|
|
"""Run full DNS recon on multiple domains in one dnsx invocation."""
|
|
if not domains:
|
|
return []
|
|
|
|
result = self.client.recon(domains)
|
|
|
|
if not result.success and not result.has_output:
|
|
logger.warning(f"dnsx batch recon failed: {result.error}")
|
|
return []
|
|
|
|
return parse_dnsx_jsonl(result.json_records)
|
|
|
|
def do_query(
|
|
self,
|
|
domain: str,
|
|
record_types: list[str],
|
|
) -> DnsxReconResponse:
|
|
"""Run targeted DNS query for specific record types."""
|
|
result = self.client.query(domain, record_types=record_types)
|
|
|
|
if not result.success and not result.has_output:
|
|
msg = f"dnsx query failed for {domain}: {result.error or 'no output'}"
|
|
raise DnsxLookupError(msg)
|
|
|
|
parsed = parse_dnsx_jsonl(result.json_records)
|
|
if not parsed:
|
|
msg = f"dnsx returned no parseable results for {domain}"
|
|
raise DnsxLookupError(msg)
|
|
|
|
return parsed[0]
|
|
|
|
def do_resolve(self, domains: list[str]) -> list[str]:
|
|
"""Check which domains resolve (have active DNS)."""
|
|
result = self.client.resolve(domains)
|
|
return result.output_lines
|
|
|
|
def lookup_and_store(
|
|
self,
|
|
session: Session,
|
|
target_id: uuid.UUID,
|
|
domain: str,
|
|
) -> DnsLookup:
|
|
"""Run DNS recon and store results linked to a target."""
|
|
response = self.do_recon(domain)
|
|
return self.store_results(session, target_id, response)
|
|
|
|
def store_results(
|
|
self,
|
|
session: Session,
|
|
target_id: uuid.UUID,
|
|
response: DnsxReconResponse,
|
|
) -> DnsLookup:
|
|
"""Store DNS recon results, replacing any existing DnsLookup for this target."""
|
|
now = utc_now()
|
|
|
|
target = session.get(Target, target_id)
|
|
if target and target.dns_lookup_id:
|
|
existing = session.get(DnsLookup, target.dns_lookup_id)
|
|
if existing:
|
|
session.delete(existing)
|
|
session.flush()
|
|
target.dns_lookup_id = None
|
|
|
|
lookup = DnsLookup(
|
|
host=response.host,
|
|
status_code=response.status_code,
|
|
cdn=response.cdn,
|
|
cdn_name=response.cdn_name,
|
|
parsed_data=response.model_dump(mode="json"),
|
|
queried_at=now,
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
session.add(lookup)
|
|
session.flush()
|
|
|
|
db_record_dicts = response.to_db_records()
|
|
for rec_dict in db_record_dicts:
|
|
dns_record = DnsRecord(
|
|
dns_lookup_id=lookup.id,
|
|
target_id=target_id,
|
|
record_type=DnsRecordType(rec_dict["record_type"]),
|
|
value=rec_dict.get("value", ""),
|
|
priority=rec_dict.get("priority"),
|
|
weight=rec_dict.get("weight"),
|
|
port=rec_dict.get("port"),
|
|
soa_email=rec_dict.get("soa_email"),
|
|
soa_serial=rec_dict.get("soa_serial"),
|
|
soa_refresh=rec_dict.get("soa_refresh"),
|
|
soa_retry=rec_dict.get("soa_retry"),
|
|
soa_expire=rec_dict.get("soa_expire"),
|
|
soa_minttl=rec_dict.get("soa_minttl"),
|
|
caa_tag=rec_dict.get("caa_tag"),
|
|
caa_flag=rec_dict.get("caa_flag"),
|
|
created_at=now,
|
|
)
|
|
session.add(dns_record)
|
|
|
|
session.commit()
|
|
session.refresh(lookup)
|
|
return lookup
|
|
|
|
def get_lookup_sync(
|
|
self, session: Session, target_id: uuid.UUID
|
|
) -> DnsLookup | None:
|
|
"""Get existing DNS lookup for a target (sync)."""
|
|
target = session.get(Target, target_id)
|
|
if target and target.dns_lookup_id:
|
|
return session.get(DnsLookup, target.dns_lookup_id)
|
|
return None
|
|
|
|
def find_targets_by_record(
|
|
self,
|
|
session: Session,
|
|
record_type: DnsRecordType,
|
|
value: str,
|
|
limit: int = MAX_CORRELATION_TARGETS,
|
|
) -> list[DnsRecord]:
|
|
"""Find all DNS records matching a type and value (rows include target_id for joining)."""
|
|
result = session.execute(
|
|
select(DnsRecord)
|
|
.where(DnsRecord.record_type == record_type)
|
|
.where(DnsRecord.value == value)
|
|
.limit(limit)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
def find_targets_sharing_ns(
|
|
self,
|
|
session: Session,
|
|
nameserver: str,
|
|
include_shared: bool = False,
|
|
limit: int = MAX_CORRELATION_TARGETS,
|
|
) -> list[uuid.UUID]:
|
|
"""Target IDs sharing a nameserver (shared NS skipped unless include_shared)."""
|
|
if not include_shared and is_shared_nameserver(nameserver):
|
|
return []
|
|
result = session.execute(
|
|
select(DnsRecord.target_id)
|
|
.where(DnsRecord.record_type == DnsRecordType.NS)
|
|
.where(DnsRecord.value == nameserver)
|
|
.distinct()
|
|
.limit(limit)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
def find_targets_sharing_mx(
|
|
self,
|
|
session: Session,
|
|
mx_host: str,
|
|
limit: int = MAX_CORRELATION_TARGETS,
|
|
) -> list[uuid.UUID]:
|
|
"""Find all target IDs that share a given MX host."""
|
|
result = session.execute(
|
|
select(DnsRecord.target_id)
|
|
.where(DnsRecord.record_type == DnsRecordType.MX)
|
|
.where(DnsRecord.value == mx_host)
|
|
.distinct()
|
|
.limit(limit)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
def find_targets_sharing_ip(
|
|
self,
|
|
session: Session,
|
|
ip: str,
|
|
exclude_cdn: bool = True,
|
|
limit: int = MAX_CORRELATION_TARGETS,
|
|
) -> list[uuid.UUID]:
|
|
"""Target IDs resolving to an IP (CDN-fronted excluded unless exclude_cdn=False)."""
|
|
query = (
|
|
select(DnsRecord.target_id)
|
|
.where(DnsRecord.record_type.in_([DnsRecordType.A, DnsRecordType.AAAA]))
|
|
.where(DnsRecord.value == ip)
|
|
)
|
|
if exclude_cdn:
|
|
query = query.join(
|
|
DnsLookup, DnsLookup.id == DnsRecord.dns_lookup_id
|
|
).where(DnsLookup.cdn.is_(False))
|
|
result = session.execute(query.distinct().limit(limit))
|
|
return list(result.scalars().all())
|
|
|
|
def get_record_value_counts(
|
|
self,
|
|
session: Session,
|
|
record_type: DnsRecordType,
|
|
limit: int = MAX_CORRELATION_TARGETS,
|
|
) -> list[tuple[str, int]]:
|
|
"""Get value counts for a record type, ordered by frequency."""
|
|
result = session.execute(
|
|
select(DnsRecord.value, func.count(DnsRecord.target_id).label("count"))
|
|
.where(DnsRecord.record_type == record_type)
|
|
.group_by(DnsRecord.value)
|
|
.order_by(func.count(DnsRecord.target_id).desc())
|
|
.limit(limit)
|
|
)
|
|
return list(result.all())
|
|
|
|
@staticmethod
|
|
def to_lookup_summary(lookup: DnsLookup) -> DnsLookupSummary:
|
|
"""Convert a DnsLookup DB record to a summary for API responses."""
|
|
record_counts: dict[str, int] = {}
|
|
if lookup.records:
|
|
counts = Counter(r.record_type.value for r in lookup.records)
|
|
record_counts = dict(counts)
|
|
|
|
return DnsLookupSummary(
|
|
id=lookup.id,
|
|
host=lookup.host,
|
|
status_code=lookup.status_code,
|
|
cdn=lookup.cdn,
|
|
cdn_name=lookup.cdn_name,
|
|
queried_at=lookup.queried_at,
|
|
record_counts=record_counts,
|
|
)
|