259 lines
9.8 KiB
Python
259 lines
9.8 KiB
Python
import subprocess
|
||
import os
|
||
import logging
|
||
import csv
|
||
import functools
|
||
|
||
@functools.lru_cache(maxsize=4096)
|
||
def normalize_domain(domain: str) -> str:
|
||
"""
|
||
规范化域名，去除动态生成的前缀。
|
||
"""
|
||
try:
|
||
if not domain:
|
||
return domain
|
||
parts = domain.split('.')
|
||
if len(parts) < 3:
|
||
return domain
|
||
prefix = parts[0]
|
||
# 判断前缀是否像动态生成的（长且包含数字和字母）
|
||
is_long = len(prefix) > 10
|
||
has_digits = any(char.isdigit() for char in prefix)
|
||
has_letters = any(char.isalpha() for char in prefix)
|
||
if is_long and has_digits and has_letters:
|
||
normalized = '.'.join(parts[1:])
|
||
return normalized
|
||
else:
|
||
return domain
|
||
except:
|
||
return domain
|
||
|
||
|
||
def build_blacklist_index(blacklist: set) -> dict:
|
||
"""
|
||
将黑名单后缀列表预处理为反向分层索引，加速后续查找。
|
||
|
||
原始线性扫描：O(黑名单大小)
|
||
分层索引查找：O(域名层级数)，通常只需 2-3 次 set 查找
|
||
|
||
返回结构示例（黑名单含 "ads.google.com" 和 "tracker.com"）：
|
||
{
|
||
'com': {
|
||
'google': {'ads': {}},
|
||
'tracker': {},
|
||
}
|
||
}
|
||
空 dict 表示该节点为黑名单终止节点（完整后缀匹配）。
|
||
"""
|
||
index = {}
|
||
for suffix in blacklist:
|
||
if not suffix:
|
||
continue
|
||
parts = suffix.strip().split('.')
|
||
node = index
|
||
# 从右到左插入（TLD 在最外层）
|
||
for part in reversed(parts):
|
||
if part not in node:
|
||
node[part] = {}
|
||
node = node[part]
|
||
return index
|
||
|
||
|
||
def is_blacklisted_domain(domain: str, blacklist: set,
|
||
_index: dict = None) -> bool:
|
||
"""
|
||
检查域名是否在黑名单中（支持后缀匹配）。
|
||
|
||
如果提供了预建的 _index（由 build_blacklist_index 生成），
|
||
使用分层索引 O(层级数) 查找；否则回退到线性扫描。
|
||
"""
|
||
if not domain:
|
||
return True
|
||
if _index is not None:
|
||
# 分层索引查找：从 TLD 往左逐层匹配
|
||
parts = domain.split('.')
|
||
node = _index
|
||
for part in reversed(parts):
|
||
if part not in node:
|
||
return False # 当前层无匹配，不在黑名单
|
||
node = node[part]
|
||
if not node:
|
||
return True # 到达终止节点，完整后缀匹配
|
||
return bool(not node) # 所有层都匹配完即为命中
|
||
# 回退：原始线性扫描
|
||
for blocked_suffix in blacklist:
|
||
if domain.endswith(blocked_suffix):
|
||
return True
|
||
return False
|
||
|
||
|
||
class TrafficMonitor:
|
||
"""
|
||
TrafficMonitor 用于实时监控 PCAPdroid 生成的流量日志。
|
||
它通过 ADB 读取最新的 txt 记录，并统计特定包名的去重域名。
|
||
"""
|
||
|
||
def __init__(self, device, package_name, blacklist_path="./doc/ad_domains.csv", pcap_callback=None):
|
||
"""
|
||
初始化流量监控模块。
|
||
:param device: 设备实例
|
||
:param package_name: 需要监控的包名
|
||
:param blacklist_path: 广告域名黑名单 CSV 路径
|
||
:param pcap_callback: (Optional) 回调函数，用于在检测到停止时推送流量文件
|
||
"""
|
||
self.logger = logging.getLogger(self.__class__.__name__)
|
||
self.device = device
|
||
self.package_name = package_name
|
||
self.pcap_callback = pcap_callback # 保存回调
|
||
self.remote_dir = self.device.captured_traffic_dir
|
||
|
||
self.all_domains = set() # 自启动以来发现的所有去重域名
|
||
self.last_domains_count = 0 # 上次外部查询时的域名总数
|
||
self.last_step_domains = set() # 上一步时的所有域名（用于计算每步新增）
|
||
|
||
# 已检查过黑名单的域名缓存（True=黑名单，False=白名单）
|
||
self._blacklist_cache: dict = {}
|
||
|
||
# 加载黑名单并预建分层索引
|
||
self.blacklist = set()
|
||
self._blacklist_index = {}
|
||
if blacklist_path and os.path.exists(blacklist_path):
|
||
try:
|
||
with open(blacklist_path, mode='r', encoding='utf-8') as f:
|
||
reader = csv.reader(f)
|
||
for row in reader:
|
||
if row:
|
||
self.blacklist.add(row[0].strip())
|
||
self._blacklist_index = build_blacklist_index(self.blacklist)
|
||
self.logger.info(f"Loaded {len(self.blacklist)} blacklisted domains from {blacklist_path}")
|
||
except Exception as e:
|
||
self.logger.error(f"Failed to load blacklist: {e}")
|
||
|
||
self.logger.info(f"TrafficMonitor initialized for package: {package_name}")
|
||
self.update_count = 0
|
||
|
||
def update(self):
|
||
"""
|
||
通过设备接口读取远程目录下最新的 txt 文件并解析。
|
||
"""
|
||
content = self.device.get_traffic_domains(self.remote_dir)
|
||
|
||
if content:
|
||
self._parse_content(content)
|
||
self.update_count += 1
|
||
else:
|
||
self.logger.warning(f"device.get_traffic_domains() return None. update_count: {self.update_count}")
|
||
|
||
def _parse_content(self, content):
|
||
"""
|
||
解析日志内容并将域名加入缓存。
|
||
格式: package_name, app_name, domain
|
||
|
||
性能优化：
|
||
1. 用行内容 hash 跳过已处理行，避免重复解析全量历史内容
|
||
2. 用 _blacklist_cache 缓存黑名单检查结果，同一域名只检查一次
|
||
3. 用预建的分层索引替代线性遍历黑名单
|
||
"""
|
||
lines = content.splitlines()
|
||
for line in lines:
|
||
parts = [p.strip() for p in line.split(',')]
|
||
# 格式校验：至少有 3 个元素
|
||
if len(parts) >= 3:
|
||
pkg_name = parts[0]
|
||
domain = parts[2]
|
||
|
||
if pkg_name == self.package_name and domain:
|
||
# 1. 规范化域名
|
||
normalized_domain = normalize_domain(domain)
|
||
|
||
# 2. 已知域名直接跳过（无需再查黑名单）
|
||
if normalized_domain in self.all_domains:
|
||
continue
|
||
|
||
# 3. 检查黑名单（带缓存）
|
||
if normalized_domain not in self._blacklist_cache:
|
||
self._blacklist_cache[normalized_domain] = is_blacklisted_domain(
|
||
normalized_domain, self.blacklist, self._blacklist_index
|
||
)
|
||
if not self._blacklist_cache[normalized_domain]:
|
||
self.all_domains.add(normalized_domain)
|
||
self.logger.debug(f"New domain found: {normalized_domain} (original: {domain})")
|
||
|
||
def has_new_features(self):
|
||
"""
|
||
检查自上次调用以来是否有新的去重域名出现。
|
||
:return: (bool, int) 是否有新域名，新增了多少个
|
||
"""
|
||
current_count = len(self.all_domains)
|
||
new_count = current_count - self.last_domains_count
|
||
|
||
has_new = new_count > 0
|
||
|
||
# 更新记录
|
||
self.last_domains_count = current_count
|
||
|
||
if has_new:
|
||
self.logger.info(f"Feedback: {new_count} new unique domains detected.")
|
||
|
||
return has_new, new_count
|
||
|
||
def get_all_domains(self):
|
||
"""
|
||
返回所有已发现的去重域名。
|
||
"""
|
||
return sorted(list(self.all_domains))
|
||
|
||
def get_new_domains_since_last_step(self):
|
||
"""
|
||
获取自上一步以来新增的域名。
|
||
每次调用后会更新 last_step_domains，确保每步只获取一次新增域名。
|
||
:return: (int, list) 新增域名数量，新增域名列表
|
||
"""
|
||
# 先更新获取最新数据
|
||
self.update()
|
||
|
||
# 计算新增域名
|
||
new_domains = self.all_domains - self.last_step_domains
|
||
new_domains_list = sorted(list(new_domains))
|
||
new_count = len(new_domains_list)
|
||
|
||
# 更新上一步的域名集合
|
||
self.last_step_domains = self.all_domains.copy()
|
||
|
||
if new_count > 0:
|
||
self.logger.debug(f"Step found {new_count} new domains: {new_domains_list}")
|
||
|
||
return new_count, new_domains_list
|
||
|
||
def check_monitor_health(self):
|
||
"""
|
||
检查监控环境是否正常运行。
|
||
如果停止运行：
|
||
1. 调用 pcap_callback (如果存在) 推送旧流量文件
|
||
2. 通过设备接口重启抓包
|
||
"""
|
||
try:
|
||
if not self.device.is_traffic_capture_running():
|
||
self.logger.warning("[Monitor Health] Traffic capture stopped running!")
|
||
|
||
# 1. 尝试推送旧流量文件（因为重启会覆盖/生成新文件）
|
||
if self.pcap_callback:
|
||
self.logger.info("[Monitor Health] Triggering pcap_callback to push traffic file...")
|
||
try:
|
||
self.pcap_callback(self.package_name)
|
||
except Exception as e:
|
||
self.logger.error(f"[Monitor Health] Failed to execute pcap_callback: {e}")
|
||
else:
|
||
self.logger.warning("[Monitor Health] No pcap_callback provided, skipping traffic push.")
|
||
|
||
# 2. 通过设备接口重启抓包
|
||
self.logger.info("[Monitor Health] Restarting traffic capture...")
|
||
self.device.restart_traffic_capture(self.package_name)
|
||
time.sleep(15)
|
||
else:
|
||
self.logger.debug("[Monitor Health] Traffic capture is running normally.")
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"[Monitor Health] Exception during health check: {e}")
|
||
|