├── README.md
└── screen_get.py
/README.md:
--------------------------------------------------------------------------------
1 | 1.安装依赖 pip install selenium webdriver-manager tqdm
2 |
3 | 2.在py脚本所在文件夹下，创建targets.txt文件，运行脚本即可。
4 |
5 | 说明：
6 |
7 | 感谢夜组安全公众号的转发，因为第一版的效果不是很好。目前正在重构相关的截图逻辑，加入进度条显示。优化截图速度特别是面对资产多的情况。
8 | 新版本已上线！！！！，放在Releases，有相关问题提issue
9 |
10 | (1)该脚本旨在攻防中，批量的请求url，获取页面截图，以便初步判断脆弱资产。
11 |
12 | (2)脚本运行后，会在py文件所在的文件夹下创建screenshots文件夹，并把截图保存到该文件夹下，最后会生成html文档，以便快速筛选可能脆弱的资产
13 |
14 | 运行截图：
15 |
16 |
17 |
18 | 
19 |
20 |
21 | 1.2版本已更新实现src资产的妙级截图
22 |
23 | 
24 |
--------------------------------------------------------------------------------
/screen_get.py:
--------------------------------------------------------------------------------
1 | import os
2 | import re
3 | import time
4 | import random # 导入 random 用于生成随机数
5 | import concurrent.futures
6 | from tqdm import tqdm # 导入 tqdm 用于进度条
7 | from queue import PriorityQueue # 使用优先队列来根据响应时间优先处理
8 | from selenium import webdriver
9 | from selenium.webdriver.chrome.service import Service
10 | from selenium.webdriver.chrome.options import Options
11 | from webdriver_manager.chrome import ChromeDriverManager
12 |
13 | # 初始化 Chrome 浏览器的设置
14 | chrome_options = Options()
15 | chrome_options.add_argument("--headless") # 无头模式（不显示浏览器窗口）
16 | chrome_options.add_argument("--no-sandbox")
17 | chrome_options.add_argument("--disable-dev-shm-usage")
18 |
19 |
20 | # 自动下载 ChromeDriver
21 | def get_driver():
22 | return webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=chrome_options)
23 |
24 |
25 | # 创建一个用于存储截图的文件夹
26 | output_folder = "screenshots"
27 | os.makedirs(output_folder, exist_ok=True)
28 |
29 | # 创建 HTML 文件并写入初始内容
30 | html_file = "screenshots_report.html"
31 | with open(html_file, "w", encoding="utf-8") as f:
32 | f.write("URL Screenshot Report\n")
33 | f.write("\n")
34 | f.write("| URL | Screenshot |
|---|
\n") # 添加表格头部
35 |
36 |
37 | # 使用随机数和时间戳生成唯一文件名
38 | def generate_filename(url):
39 | timestamp = time.strftime("%Y%m%d-%H%M%S")
40 | random_number = random.randint(1000, 9999) # 生成 1000 到 9999 的随机数
41 | file_name = f"screenshot_{random_number}_{timestamp}.png"
42 | return file_name
43 |
44 |
45 | # 定义访问和截图的函数，并测量访问时间
46 | def capture_screenshot(url, priority_queue):
47 | driver = get_driver() # 每个线程单独初始化一个 driver
48 | start_time = time.time()
49 | try:
50 | driver.get(url)
51 | print(f"正在访问 {url}...")
52 |
53 | # 生成截图文件的名称
54 | file_name = generate_filename(url)
55 | file_path = os.path.join(output_folder, file_name)
56 |
57 | # 截图并保存到指定路径
58 | driver.save_screenshot(file_path)
59 | print(f"已保存截图：{file_path}")
60 |
61 | # 截图完成后立即写入 HTML 文件
62 | with open(html_file, "a", encoding="utf-8") as f:
63 | f.write(f"| {url} |  |
\n")
64 |
65 | # 记录完成后的响应时间
66 | end_time = time.time()
67 | response_time = end_time - start_time
68 | priority_queue.put((response_time, url)) # 将 URL 和响应时间放入优先队列
69 |
70 | return (url, file_path)
71 | finally:
72 | driver.quit()
73 |
74 |
75 | # 使用优先队列动态调整优先级
76 | def prioritize_and_process_urls(urls):
77 | priority_queue = PriorityQueue()
78 |
79 | # 将初始 URL 放入优先队列，初始优先级为 0
80 | for url in urls:
81 | priority_queue.put((0, url))
82 |
83 | # 使用多线程加速访问和截图
84 | max_threads = 10
85 | with concurrent.futures.ThreadPoolExecutor(max_threads) as executor:
86 | futures = []
87 | progress = tqdm(total=len(urls), desc="Processing URLs")
88 |
89 | while not priority_queue.empty():
90 | # 获取优先级最高的 URL（即响应时间最短的）
91 | priority, url = priority_queue.get()
92 |
93 | # 提交任务并将优先队列传递给 capture_screenshot
94 | future = executor.submit(capture_screenshot, url, priority_queue)
95 | futures.append(future)
96 |
97 | # 更新进度条
98 | progress.update(1)
99 |
100 | # 等待所有任务完成
101 | concurrent.futures.wait(futures)
102 | progress.close()
103 |
104 |
105 | # 从 targets.txt 文件中读取 URL 列表
106 | with open("targets.txt", "r", encoding="utf-8") as f:
107 | urls = [line.strip() for line in f if line.strip()] # 去除空行和多余空格
108 |
109 | # 优先处理并开始截图
110 | prioritize_and_process_urls(urls)
111 |
112 | # 最后关闭 HTML 文件的标签
113 | with open(html_file, "a", encoding="utf-8") as f:
114 | f.write("
\n")
115 |
116 | print(f"报告已生成：{html_file}")
117 |
--------------------------------------------------------------------------------