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")