# -*- encoding=utf8 -*- import subprocess import os import shutil import socket import uuid class DataManager: """负责文件的拉取及推送到局域网""" DEFAULT_NET_USE_USER = "tplink" DEFAULT_NET_USE_PASSWORD = "admin" DEFAULT_NET_USE_TIMEOUT_SECONDS = 15 def __init__(self, config: dict): self.config = config self.worker_id = self._generate_worker_id() self._authorized_shares = set() print(f"[DataManager] Worker ID: {self.worker_id}") def _get_ip_address(self): """获取本机IP地址（优先获取192.168开头的IP）""" try: hostname = socket.gethostname() ip_addresses = socket.gethostbyname_ex(hostname)[2] for ip in ip_addresses: if ip.startswith('192.168'): return ip if ip_addresses: return ip_addresses[0] return '127.0.0.1' except Exception as e: print(f"[DataManager] 获取IP地址失败: {e}") return '127.0.0.1' def _get_mac_address(self) -> str: """获取本机Mac地址""" try: mac = uuid.getnode() mac_address = ':'.join(['{:02x}'.format((mac >> elements) & 0xff) for elements in range(0, 8*6, 8)][::-1]) return mac_address.upper() except Exception as e: print(f"[DataManager] 获取MAC地址失败: {e}") return '00:00:00:00:00:00' def _generate_worker_id(self) -> str: """生成Worker ID (MAC_IP格式)""" ip = self._get_ip_address() mac = self._get_mac_address() # 将MAC地址中的冒号替换为横线，使其适合作为目录名 mac_safe = mac.replace(':', '-') return f"{mac_safe}_{ip}" def _run_adb(self, args): return subprocess.run(["adb"] + args, capture_output=True, text=True) def _extract_unc_share_root(self, path): """从 UNC 路径提取共享根路径，如 \\host\share。""" if not path: return None normalized = str(path).replace("/", "\\") if not normalized.startswith("\\\\"): return None parts = [part for part in normalized.lstrip("\\").split("\\") if part] if len(parts) < 2: return None return f"\\\\{parts[0]}\\{parts[1]}" def _share_exists(self, share_root): try: return os.path.exists(share_root) except OSError: return False def _ensure_share_access(self, path): """访问共享目录前，按需执行 net use 认证。""" if os.name != "nt": return True share_root = self._extract_unc_share_root(path) if not share_root: return True if share_root in self._authorized_shares and self._share_exists(share_root): return True if self._share_exists(share_root): self._authorized_shares.add(share_root) return True print(f"[Data] 共享目录未就绪，尝试重新认证: {share_root}") try: subprocess.run( ["net", "use", share_root, "/delete", "/y"], capture_output=True, text=True, timeout=self.DEFAULT_NET_USE_TIMEOUT_SECONDS, ) res = subprocess.run( [ "net", "use", share_root, self.DEFAULT_NET_USE_PASSWORD, f"/user:{self.DEFAULT_NET_USE_USER}", ], capture_output=True, text=True, timeout=self.DEFAULT_NET_USE_TIMEOUT_SECONDS, ) except Exception as e: print(f"[Data] 执行 net use 失败: {e}") return False if res.returncode != 0: error_msg = res.stderr.strip() or res.stdout.strip() or f"return code {res.returncode}" print(f"[Data] 共享目录认证失败 {share_root}: {error_msg}") return False if not self._share_exists(share_root): print(f"[Data] 共享目录认证后仍不可访问: {share_root}") return False self._authorized_shares.add(share_root) print(f"[Data] 共享目录认证成功: {share_root}") return True def clear_remote_traffic_dir(self): """测试开始前清空设备上的流量文件夹""" remote_dir = "/sdcard/Download/PCAPdroid" print(f"[Data] 正在清空设备流量目录: {remote_dir}") self._run_adb(["shell", f"rm -rf {remote_dir}/*"]) def sync_all_traffic_data(self): """全部测试完成后，统一从设备拉取整个流量文件夹并推送到局域网""" print(f"\n[Data] 开始统一拉取所有流量数据...") remote_dir = "/sdcard/Download/PCAPdroid" # 1. 确保目标局域网目录存在 (使用 worker_id: MAC_IP) target_share = os.path.join(self.config['TRAFFIC_DATA_SHARE'], self.worker_id) if not self._ensure_share_access(target_share): return False try: if not os.path.exists(target_share): os.makedirs(target_share, exist_ok=True) except Exception as e: print(f"[Data] 无法创建局域网目录 {target_share}: {e}") return False # 2. 批量拉取到局域网目录 print(f"[Data] 正在将设备文件夹 {remote_dir} 拉取至 {target_share}") res = self._run_adb(["pull", remote_dir, target_share]) if res.returncode == 0: print(f"[Data] 流量数据同步成功!") return True else: print(f"[Data] 流量数据同步失败: {res.stderr}") return False def sync_latest_traffic_file(self, package_name): """将设备上最新的流量 txt 文件拉取并推送到局域网指定目录""" remote_traffic_dir = "/sdcard/Download/PCAPdroid/traffic" # 1. 查找设备上最新的 txt 文件 # 使用 ls -t 按修改时间排序，取第一个 cmd = f"ls -t {remote_traffic_dir}/*.txt | head -n 1" res = self._run_adb(["shell", cmd]) if res.returncode != 0 or not res.stdout.strip(): print(f"[Data] 未在设备上找到流量 txt 文件: {remote_traffic_dir}") return False remote_file_path = res.stdout.strip() file_name = os.path.basename(remote_file_path) # 1.5 确认文件内包含当前应用的流量数据（每行第一列为包名） check_cmd = f"grep -c '^{package_name}' {remote_file_path} || true" check_res = self._run_adb(["shell", check_cmd]) if check_res.returncode == 0 and check_res.stdout.strip().isdigit(): count = int(check_res.stdout.strip()) if count == 0: print(f"[Data] 流量文件中未找到 {package_name} 的数据，跳过同步") return False else: print(f"[Data] 无法检查流量文件内容，跳过同步") return False # 2. 准备局域网目标路径 (TRAFFIC_DATA_SHARE / worker_id(MAC_IP) / package_name) target_dir = os.path.join(self.config['TRAFFIC_DATA_SHARE'], self.worker_id, package_name) if not self._ensure_share_access(target_dir): return False try: if not os.path.exists(target_dir): os.makedirs(target_dir, exist_ok=True) except Exception as e: print(f"[Data] 无法创建局域网目录 {target_dir}: {e}") return False # 3. 拉取文件 target_file_path = os.path.join(target_dir, file_name) print(f"[Data] 正在同步最新流量文件: {remote_file_path} -> {target_file_path}") res = self._run_adb(["pull", remote_file_path, target_file_path]) if res.returncode == 0: print(f"[Data] 单个流量文件同步成功!") return True else: print(f"[Data] 单个流量文件同步失败: {res.stderr}") return False def sync_traversal_log(self, local_output_dir): """将 DroidBot 产生的本地 output 文件夹复制到局域网""" if not local_output_dir or not os.path.exists(local_output_dir): print("[Data] 未找到有效的输出文件夹，跳过同步") return False # 使用 worker_id (MAC_IP) 作为目录层级 folder_name = os.path.basename(local_output_dir) target_parent = os.path.join(self.config['TRAVERSAL_LOG_SHARE'], self.worker_id) target_dir = os.path.join(target_parent, folder_name) if not self._ensure_share_access(target_parent): return False try: if not os.path.exists(target_parent): os.makedirs(target_parent, exist_ok=True) print(f"[Data] 正在将日志复制至局域网: {target_dir}") shutil.copytree(local_output_dir, target_dir, dirs_exist_ok=True) return True except Exception as e: print(f"[Data] 复制日志文件夹失败: {e}") return False def sync_log_file(self, log_filepath, target_filename=None): """将 batch_run 的日志文件复制到局域网 Args: log_filepath: 日志文件的完整路径 target_filename: 可选，指定目标文件名（不包含路径） """ if not log_filepath or not os.path.exists(log_filepath): print("[Data] 未找到有效的日志文件，跳过同步") return False # 使用 worker_id (MAC_IP) 作为目录层级 file_name = target_filename if target_filename else os.path.basename(log_filepath) target_dir = os.path.join(self.config.get('TRAVERSAL_LOG_SHARE'), self.worker_id) target_file_path = os.path.join(target_dir, file_name) if not self._ensure_share_access(target_dir): return False try: if not os.path.exists(target_dir): os.makedirs(target_dir, exist_ok=True) print(f"[Data] 正在将日志文件复制至局域网: {target_file_path}") shutil.copy2(log_filepath, target_file_path) print(f"[Data] 日志文件同步成功!") return True except Exception as e: print(f"[Data] 复制日志文件失败: {e}") return False