sehoon0519 / webshot2.py

import argparse import asyncio import csv import json import random import re from datetime import datetime from pathlib import Path import xml.etree.ElementTree as ET from typing import List, Tuple, Dict, Optional, Set from shutil import copyfile from playwright.async_api import async_playwright, TimeoutError as PlaywrightTimeoutError # Defaults/assumptions DEFAULT_CONCURRENCY = 6 DEFAULT_TIMEOUT_MS = 10_000 # per navigation JITTER_MIN_S = 0.1 JITTER_MAX_S = 0.3 # 추가: 페이지 이동 후 스크린샷까지 대기 시간(ms) POST_NAV_WAIT_MS = 5_000 CSV_FIELDS = [ "timestamp", "ip", "port", "scheme_attempted", "outcome", # success | timeout | error "http_status", "final_url", "page_title", "error_message", ] def parse_nmap_xml(xml_path: Path, pairs: Set[Tuple[str, int]]) -> None: """ Parse one Nmap XML file, updating 'pairs' set of (ip, port) for open TCP ports. """ tree = ET.parse(xml_path) root = tree.getroot() for host in root.findall(".//host"): ip = None for addr in host.findall("address"): if addr.get("addrtype") == "ipv4": ip = addr.get("addr") break if not ip: continue for port_el in host.findall(".//ports/port"): if port_el.get("protocol") != "tcp": continue state_el = port_el.find("state") if state_el is None or state_el.get("state") != "open": continue portid = port_el.get("portid") if not portid or not portid.isdigit(): continue pairs.add((ip, int(portid))) def parse_nmap_text(txt_path: Path, pairs: Set[Tuple[str, int]]) -> None: """ Parse Nmap normal text output and update 'pairs' set with (ip, open_tcp_port). Handles host lines like: - Nmap scan report for 103.139.84.64 - Nmap scan report for ems1.example.org (203.0.113.10) And port lines under a host block like: - 80/tcp open http - 443/tcp open https """ host_re = re.compile(r"^Nmap scan report for\s+(?P<host>.+?)(?:\s*\((?P<ip>\d+\.\d+\.\d+\.\d+)\))?\s*$") open_port_re = re.compile(r"^(?P<port>\d{1,5})/tcp\s+open\b", re.IGNORECASE) ip_lit_re = re.compile(r"^\d+\.\d+\.\d+\.\d+$") current_ip: Optional[str] = None in_ports_section = False with txt_path.open("r", encoding="utf-8", errors="ignore") as f: for raw in f: line = raw.rstrip("\n") m_host = host_re.match(line) if m_host: # New host starts host_str = m_host.group("host").strip() ip = m_host.group("ip") if ip: current_ip = ip else: current_ip = host_str if ip_lit_re.match(host_str) else None in_ports_section = False continue if line.startswith("PORT") and "STATE" in line and "SERVICE" in line: in_ports_section = True continue if not current_ip: continue if not line.strip(): in_ports_section = False continue if in_ports_section: m_port = open_port_re.match(line.strip()) if m_port: port = int(m_port.group("port")) pairs.add((current_ip, port)) def parse_inputs(paths: List[Path]) -> List[Tuple[str, int]]: """ Parse a mix of XML and text Nmap outputs; return sorted unique (ip, port) list. """ pairs: Set[Tuple[str, int]] = set() for p in paths: suffix = p.suffix.lower() try: if suffix == ".xml": parse_nmap_xml(p, pairs) else: parse_nmap_text(p, pairs) except Exception as e: # Skip unreadable/invalid files; could log in the future print(f"[WARN] Failed to parse {p}: {e}") return sorted(pairs, key=lambda t: (t[0], t[1])) def ensure_ip_dirs(base_dir: Path, pairs: List[Tuple[str,int]]) -> None: ips = sorted({ip for ip, _ in pairs}) for ip in ips: (base_dir / ip).mkdir(parents=True, exist_ok=True) def prepare_logs(base_dir: Path) -> Tuple[Path, Path]: logs_dir = base_dir / "logs" logs_dir.mkdir(parents=True, exist_ok=True) return logs_dir / "webshot.csv", logs_dir / "webshot.json" def write_logs(csv_path: Path, json_path: Path, records: List[Dict]) -> None: # CSV: append with header if file is new is_new_csv = not csv_path.exists() with csv_path.open("a", newline="", encoding="utf-8") as fcsv: writer = csv.DictWriter(fcsv, fieldnames=CSV_FIELDS) if is_new_csv: writer.writeheader() for rec in records: writer.writerow(rec) # JSON: NDJSON (one object per line) with json_path.open("a", encoding="utf-8") as fjson: for rec in records: fjson.write(json.dumps(rec, ensure_ascii=False) + "\n") async def attempt_and_screenshot(context, url: str, timeout_ms: int, screenshot_path: Path, progress_label: str) -> Tuple[str, Optional[int], Optional[str], Optional[str], Optional[str]]: r""" Try to navigate to the given URL and always save a screenshot to screenshot_path. Returns a tuple: (outcome, status, final_url, title, error_message) - outcome: "success" | "timeout" | "error" - status: HTTP status code if any - final_url: the page.url after navigation (if any) - title: page title if retrievable - error_message: string on failure Note: Always attempts to wait an extra POST_NAV_WAIT_MS (default 5000 ms) before taking a screenshot, and also tries network idle state if possible. """ page = await context.new_page() print(f"[START] {progress_label}: {url}", flush=True) try: resp = await page.goto(url, wait_until="domcontentloaded", timeout=timeout_ms) status = resp.status if resp else None final_url = page.url title = None # 추가 대기: 네트워크 안정화 시도 후, 고정 5초 대기 try: # networkidle은 모든 네트워크가 잠잠해질 때까지 대기(환경에 따라 오래 걸릴 수 있어 Optional) await page.wait_for_load_state("networkidle", timeout=timeout_ms) except Exception: # networkidle 실패/시간초과는 무시하고 고정 대기만 수행 pass await page.wait_for_timeout(POST_NAV_WAIT_MS) try: title = await page.title() except Exception: title = None # Save screenshot of the loaded page (any status is fine) screenshot_path.parent.mkdir(parents=True, exist_ok=True) await page.screenshot(path=str(screenshot_path), full_page=True) print(f"[OK] {progress_label}: status={status} saved={screenshot_path}", flush=True) return "success", status, final_url, title, None except PlaywrightTimeoutError as e: # Try to capture whatever is visible; if not, render an error HTML and screenshot it try: screenshot_path.parent.mkdir(parents=True, exist_ok=True) await page.screenshot(path=str(screenshot_path), full_page=True) except Exception: try: await page.set_content( f""" <html><body style='font-family:Segoe UI,Tahoma,Arial,sans-serif;padding:24px;'> <h2 style='color:#b00;margin:0 0 8px'>Timeout</h2> <div>URL: {url}</div> <div>Reason: {str(e)}</div> </body></html> """ ) await page.screenshot(path=str(screenshot_path), full_page=True) except Exception: pass print(f"[TIME] {progress_label}: timeout saved={screenshot_path}", flush=True) return "timeout", None, None, None, str(e) except Exception as e: # Generic error: same fallback to ensure a screenshot try: screenshot_path.parent.mkdir(parents=True, exist_ok=True) await page.screenshot(path=str(screenshot_path), full_page=True) except Exception: try: await page.set_content( f""" <html><body style='font-family:Segoe UI,Tahoma,Arial,sans-serif;padding:24px;'> <h2 style='color:#b00;margin:0 0 8px'>Connection Error</h2> <div>URL: {url}</div> <div>Reason: {str(e)}</div> </body></html> """ ) await page.screenshot(path=str(screenshot_path), full_page=True) except Exception: pass print(f"[ERR] {progress_label}: error saved={screenshot_path}", flush=True) return "error", None, None, None, str(e) finally: try: await page.close() except Exception: pass async def process_target(ip: str, port: int, context, base_dir: Path, timeout_ms: int, sem: asyncio.Semaphore) -> List[Dict]: r""" For one (ip, port): try HTTP first; on failure/timeout, try HTTPS (ignore HTTPS cert errors). Always save screenshots per attempt under <ip>\_attempts\ and ensure final <ip>\<port>.png exists. Returns a list of 1–2 log records (one per attempt performed). """ async with sem: # Jitter to reduce burstiness await asyncio.sleep(random.uniform(JITTER_MIN_S, JITTER_MAX_S)) ts = datetime.now().isoformat(timespec="seconds") http_url = f"http://{ip}:{port}" https_url = f"https://{ip}:{port}" ip_dir = base_dir / ip final_screenshot = ip_dir / f"{port}.png" attempts_dir = ip_dir / "_attempts" attempts_dir.mkdir(parents=True, exist_ok=True) http_attempt_path = attempts_dir / f"{port}_http.png" https_attempt_path = attempts_dir / f"{port}_https.png" records: List[Dict] = [] # Attempt HTTP first outcome, status, final_url, title, err_msg = await attempt_and_screenshot( context, http_url, timeout_ms, http_attempt_path, f"{ip}:{port} HTTP" ) http_rec = { "timestamp": ts, "ip": ip, "port": port, "scheme_attempted": "http", "outcome": outcome, "http_status": status, "final_url": final_url, "page_title": title, "error_message": err_msg, } records.append(http_rec) if outcome == "success": # Promote HTTP attempt to final screenshot try: copyfile(http_attempt_path, final_screenshot) except Exception: pass print(f"[SAVE] {ip}:{port} -> {final_screenshot} (http)", flush=True) return records # HTTP failed -> Attempt HTTPS outcome2, status2, final_url2, title2, err_msg2 = await attempt_and_screenshot( context, https_url, timeout_ms, https_attempt_path, f"{ip}:{port} HTTPS" ) https_rec = { "timestamp": ts, "ip": ip, "port": port, "scheme_attempted": "https", "outcome": outcome2, "http_status": status2, "final_url": final_url2, "page_title": title2, "error_message": err_msg2, } records.append(https_rec) # Promote HTTPS attempt to final if success; otherwise keep HTTPS error as final try: src = https_attempt_path if https_attempt_path.exists() else http_attempt_path copyfile(src, final_screenshot) except Exception: pass print(f"[SAVE] {ip}:{port} -> {final_screenshot} ({'https' if outcome2=='success' else 'error'})", flush=True) return records async def run(input_paths: List[Path], concurrency: int, timeout_ms: int, base_dir: Path) -> None: csv_path, json_path = prepare_logs(base_dir) pairs = parse_inputs(input_paths) total = len(pairs) if not pairs: print("No open TCP ports found in the provided Nmap files.") return ensure_ip_dirs(base_dir, pairs) print(f"Targets: {total} open TCP ports across {len(set(ip for ip,_ in pairs))} hosts; Concurrency={concurrency}; Timeout={timeout_ms}ms", flush=True) async with async_playwright() as pw: # 번들 브라우저가 없을 경우를 대비하여, 시스템 Chrome 사용으로 바꾸고 싶다면: # browser = await pw.chromium.launch(executable_path=r"C:\Program Files\Google\Chrome\Application\chrome.exe", headless=True) browser = await pw.chromium.launch(headless=True) context = await browser.new_context(ignore_https_errors=True) sem = asyncio.Semaphore(concurrency) tasks = [ asyncio.create_task(process_target(ip, port, context, base_dir, timeout_ms, sem)) for ip, port in pairs ] results: List[Dict] = [] done = 0 for coro in asyncio.as_completed(tasks): recs = await coro results.extend(recs) done += 1 # Progress line using the first record ip = recs[0].get("ip") port = recs[0].get("port") summary = ("success" if any(r.get("outcome") == "success" for r in recs) else ("timeout" if any(r.get("outcome") == "timeout" for r in recs) else "error")) print(f"[PROG] {done}/{total} completed: {ip}:{port} => {summary}", flush=True) await context.close() await browser.close() # Write all attempt records write_logs(csv_path, json_path, results) print(f"[DONE] Logs written: {csv_path} and {json_path}", flush=True) def cleanup_final_images(base_dir: Path) -> int: r""" Remove top-level <port>.png files under each IP folder, keeping _attempts intact. Returns number of files deleted. """ deleted = 0 for ip_dir in base_dir.iterdir(): try: if not ip_dir.is_dir(): continue # Skip common system/profile directories when base_dir is a user profile root if ip_dir.name.lower() in {"application data", "appdata", "saved games", "documents", "downloads", "music", "pictures", "videos", "contacts", "links", "searches", "favorites", "onedrive"}: continue if ip_dir.name == "logs": continue attempts_dir = ip_dir / "_attempts" for item in ip_dir.iterdir(): try: if item.is_dir(): continue if item.suffix.lower() != ".png": continue # Only delete files named like "<digits>.png" (e.g., 80.png, 443.png) if re.fullmatch(r"\d{1,5}\.png", item.name): try: item.unlink() deleted += 1 except Exception: pass except PermissionError: continue except PermissionError: continue return deleted def main(): parser = argparse.ArgumentParser(description="Screenshot web UIs from Nmap results (XML or text). HTTP→HTTPS, Playwright.") parser.add_argument("inputs", nargs="*", help="Nmap output files (XML or normal text), e.g., scan1.xml scan2.xml scan3.txt") parser.add_argument("--concurrency", type=int, default=DEFAULT_CONCURRENCY, help="Max parallel targets (default: 6)") parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT_MS, help="Navigation timeout in ms (default: 10000)") parser.add_argument("--cleanup-final", action="store_true", help="Remove top-level <port>.png files, keep only _attempts images.") parser.add_argument("--base-dir", default=None, help="Directory where outputs/logs will be written (default: current working directory)") args = parser.parse_args() base_dir = Path(args.base_dir) if args.base_dir else Path.cwd() if not args.inputs and args.cleanup_final: deleted = cleanup_final_images(base_dir) print(f"[CLEANUP] Deleted {deleted} top-level port PNGs (kept _attempts)") return input_paths = [Path(p) for p in args.inputs] for p in input_paths: if not p.exists(): raise FileNotFoundError(f"File not found: {p}") asyncio.run(run(input_paths, concurrency=args.concurrency, timeout_ms=args.timeout, base_dir=base_dir)) if args.cleanup_final: deleted = cleanup_final_images(base_dir) print(f"[CLEANUP] Deleted {deleted} top-level port PNGs (kept _attempts)") if __name__ == "__main__": main()