diff --git a/pyproject.toml b/pyproject.toml index fb1fd51112..f5d25aff1c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,7 +79,7 @@ exploitgraph = [ image = [ "pytesseract", - "weasyprint", + "playwright", ] shellcode = [ diff --git a/tests/functional/test_screenshot.py b/tests/functional/test_screenshot.py index 7501455df9..1d36cbc347 100644 --- a/tests/functional/test_screenshot.py +++ b/tests/functional/test_screenshot.py @@ -1,16 +1,10 @@ -import os import logging -import pytest from thug.ThugAPI.ThugAPI import ThugAPI log = logging.getLogger("Thug") -IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" and os.getenv( - "RUNNER_OS" -) in ("Linux",) - class TestScreenshot(object): def do_perform_test(self, caplog, url, expected, type_="remote"): @@ -38,17 +32,11 @@ def do_perform_test(self, caplog, url, expected, type_="remote"): assert matches >= len(expected) - @pytest.mark.skipif( - not (IN_GITHUB_ACTIONS), reason="Test works just in Github Actions (Linux)" - ) - def _test_antifork(self, caplog): + def test_antifork(self, caplog): expected = [] self.do_perform_test(caplog, "https://buffer.antifork.org", expected) - @pytest.mark.skipif( - not (IN_GITHUB_ACTIONS), reason="Test works just in Github Actions (Linux)" - ) - def _test_invalid_ctype(self, caplog): + def test_invalid_ctype(self, caplog): expected = [] self.do_perform_test( caplog, "https://buffer.antifork.org/images/antifork.jpg", expected diff --git a/thug/Analysis/screenshot/Screenshot.py b/thug/Analysis/screenshot/Screenshot.py index 7d59a62e34..00b50f763f 100644 --- a/thug/Analysis/screenshot/Screenshot.py +++ b/thug/Analysis/screenshot/Screenshot.py @@ -1,22 +1,22 @@ -import io import logging -import bs4 try: - import weasyprint + from playwright.sync_api import sync_playwright - WEASYPRINT_MODULE = True + PLAYWRIGHT_MODULE = True except ImportError: # pragma: no cover - WEASYPRINT_MODULE = False + PLAYWRIGHT_MODULE = False + log = logging.getLogger("Thug") class Screenshot: content_types = ("text/html",) + resource_types = ("image", "stylesheet") def __init__(self): - self.enable = WEASYPRINT_MODULE + self.enable = PLAYWRIGHT_MODULE def run(self, window, url, response, ctype): if not self.enable or not log.ThugOpts.screenshot: @@ -25,26 +25,28 @@ def run(self, window, url, response, ctype): if not ctype.startswith(self.content_types): return # pragma: no cover - soup = bs4.BeautifulSoup(response.content, "html5lib") + with sync_playwright() as p: + browser = p.chromium.launch(headless=True) + page = browser.new_page() - for img in soup.find_all("img"): - src = img.get("src", None) - if not src: - continue # pragma: no cover + def block_resource_type(route): # pragma: no cover + if route.request.resource_type in self.resource_types: + route.continue_() + else: + route.abort() - norm_src = log.HTTPSession.normalize_url(window, src) - if norm_src: - img["src"] = norm_src + page.route("**/*", block_resource_type) - content = soup.prettify(formatter=None) + try: + page.set_content(response.text) - try: - html = weasyprint.HTML(string=content) - document = html.render() + # Scroll down to enable downloading lazy-loaded images and wait + # for all the resources to be loaded + page.evaluate("window.scrollTo(0, document.body.scrollHeight)") + page.wait_for_load_state("networkidle") - with io.BytesIO() as screenshot: - document.write_pdf(screenshot) - screenshot.seek(0) - log.ThugLogging.log_screenshot(url, screenshot.read()) - except Exception as e: # pragma: no cover,pylint:disable=broad-except - log.warning("[SCREENSHOT] Error: %s", str(e)) + screenshot = page.screenshot(type="png", full_page=True) + browser.close() + log.ThugLogging.log_screenshot(url, screenshot) + except Exception as e: # pragma: no cover + log.warning("[SCREENSHOT] Error: %s", str(e)) diff --git a/thug/Logging/ThugLogging.py b/thug/Logging/ThugLogging.py index 1ce8c9bbd8..a11beb4fd5 100644 --- a/thug/Logging/ThugLogging.py +++ b/thug/Logging/ThugLogging.py @@ -452,7 +452,7 @@ def log_screenshot(self, url, screenshot): @screenshot Screenshot """ dirname = os.path.join(self.baseDir, "analysis", "screenshots") - filename = f"{hashlib.sha256(screenshot).hexdigest()}.pdf" + filename = f"{hashlib.sha256(screenshot).hexdigest()}.png" self.store_content(dirname, filename, screenshot) for m in self.resolve_method("log_screenshot"): # pragma: no cover diff --git a/tox.ini b/tox.ini index e93a8f4acb..e8d96d9966 100644 --- a/tox.ini +++ b/tox.ini @@ -30,6 +30,7 @@ commands_pre = pip install --upgrade pip pip install . pip install .[test] + playwright install chromium commands = thug --version pytest --cov-report xml