From c5fe956589125be23267d01ac95e6b4756e73da8 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Thu, 10 Sep 2026 21:34:44 +0200 Subject: [PATCH 1/9] Generate screenshots using playwright --- pyproject.toml | 2 +- thug/Analysis/screenshot/Screenshot.py | 46 ++++++++++++++++++++++++-- thug/Logging/ThugLogging.py | 2 +- tox.ini | 1 + 4 files changed, 47 insertions(+), 4 deletions(-) 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/thug/Analysis/screenshot/Screenshot.py b/thug/Analysis/screenshot/Screenshot.py index 7d59a62e34..589aaa04f3 100644 --- a/thug/Analysis/screenshot/Screenshot.py +++ b/thug/Analysis/screenshot/Screenshot.py @@ -1,10 +1,17 @@ import io +import base64 import logging import bs4 try: - import weasyprint + from playwright.sync_api import sync_playwright + PLAYWRIGHT_MODULE = True +except ImportError: # pragma: no cover + PLAYWRIGHT_MODULE = False + +try: + import weasyprint WEASYPRINT_MODULE = True except ImportError: # pragma: no cover WEASYPRINT_MODULE = False @@ -12,7 +19,7 @@ log = logging.getLogger("Thug") -class Screenshot: +class OldScreenshot: content_types = ("text/html",) def __init__(self): @@ -48,3 +55,38 @@ def run(self, window, url, response, ctype): 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)) + + +class Screenshot: + content_types = ("text/html",) + resource_types = ("image", "stylesheet") + + def __init__(self): + self.enable = PLAYWRIGHT_MODULE + + def run(self, window, url, response, ctype): + if not self.enable or not log.ThugOpts.screenshot: + return + + if not ctype.startswith(self.content_types): + return # pragma: no cover + + with sync_playwright() as p: + browser = p.chromium.launch(headless=True) + page = browser.new_page() + + def block_resource_type(route): + if route.request.resource_type in self.resource_types: + route.continue_() + else: + route.abort() + + page.route("**/*", block_resource_type) + + try: + page.set_content(response.text) + screenshot = page.screenshot(type="png", full_page=True) + browser.close() + log.ThugLogging.log_screenshot(url, screenshot) + except Exception as e: + 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 From aad382539d6af575d977bbb99b3e9ac44d0fc5e3 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Thu, 10 Sep 2026 21:45:58 +0200 Subject: [PATCH 2/9] Minor linting change --- thug/Analysis/screenshot/Screenshot.py | 1 - 1 file changed, 1 deletion(-) diff --git a/thug/Analysis/screenshot/Screenshot.py b/thug/Analysis/screenshot/Screenshot.py index 589aaa04f3..b3db0f4f1f 100644 --- a/thug/Analysis/screenshot/Screenshot.py +++ b/thug/Analysis/screenshot/Screenshot.py @@ -1,5 +1,4 @@ import io -import base64 import logging import bs4 From 54c9f79d536baa12048754ae95e1bce6e89fe71f Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Thu, 10 Sep 2026 21:51:41 +0200 Subject: [PATCH 3/9] Reenable screenshot tests --- tests/functional/test_screenshot.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/functional/test_screenshot.py b/tests/functional/test_screenshot.py index 7501455df9..1a8d462831 100644 --- a/tests/functional/test_screenshot.py +++ b/tests/functional/test_screenshot.py @@ -41,14 +41,14 @@ def do_perform_test(self, caplog, url, expected, type_="remote"): @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 From e8de71f50895d886c67e7c094c9f8332bd4cda69 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Thu, 10 Sep 2026 22:05:05 +0200 Subject: [PATCH 4/9] Minor formatting change --- thug/Analysis/screenshot/Screenshot.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/thug/Analysis/screenshot/Screenshot.py b/thug/Analysis/screenshot/Screenshot.py index b3db0f4f1f..b9f1314bbb 100644 --- a/thug/Analysis/screenshot/Screenshot.py +++ b/thug/Analysis/screenshot/Screenshot.py @@ -4,6 +4,7 @@ try: from playwright.sync_api import sync_playwright + PLAYWRIGHT_MODULE = True except ImportError: # pragma: no cover PLAYWRIGHT_MODULE = False @@ -11,6 +12,7 @@ try: import weasyprint + WEASYPRINT_MODULE = True except ImportError: # pragma: no cover WEASYPRINT_MODULE = False From 1ef0a072e29473dc1a347c1ef6830a42ab27a785 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Fri, 11 Sep 2026 10:18:06 +0200 Subject: [PATCH 5/9] Screenshot code cleanup --- thug/Analysis/screenshot/Screenshot.py | 49 +------------------------- 1 file changed, 1 insertion(+), 48 deletions(-) diff --git a/thug/Analysis/screenshot/Screenshot.py b/thug/Analysis/screenshot/Screenshot.py index b9f1314bbb..5701d1cf65 100644 --- a/thug/Analysis/screenshot/Screenshot.py +++ b/thug/Analysis/screenshot/Screenshot.py @@ -1,6 +1,4 @@ -import io import logging -import bs4 try: from playwright.sync_api import sync_playwright @@ -10,54 +8,9 @@ PLAYWRIGHT_MODULE = False -try: - import weasyprint - - WEASYPRINT_MODULE = True -except ImportError: # pragma: no cover - WEASYPRINT_MODULE = False - log = logging.getLogger("Thug") -class OldScreenshot: - content_types = ("text/html",) - - def __init__(self): - self.enable = WEASYPRINT_MODULE - - def run(self, window, url, response, ctype): - if not self.enable or not log.ThugOpts.screenshot: - return - - if not ctype.startswith(self.content_types): - return # pragma: no cover - - soup = bs4.BeautifulSoup(response.content, "html5lib") - - for img in soup.find_all("img"): - src = img.get("src", None) - if not src: - continue # pragma: no cover - - norm_src = log.HTTPSession.normalize_url(window, src) - if norm_src: - img["src"] = norm_src - - content = soup.prettify(formatter=None) - - try: - html = weasyprint.HTML(string=content) - document = html.render() - - 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)) - - class Screenshot: content_types = ("text/html",) resource_types = ("image", "stylesheet") @@ -70,7 +23,7 @@ def run(self, window, url, response, ctype): return if not ctype.startswith(self.content_types): - return # pragma: no cover + return with sync_playwright() as p: browser = p.chromium.launch(headless=True) From f2386bb7ca1f77ee9055fe9200fe85322cc3f9f1 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Fri, 11 Sep 2026 11:36:59 +0200 Subject: [PATCH 6/9] Minor improvements --- thug/Analysis/screenshot/Screenshot.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/thug/Analysis/screenshot/Screenshot.py b/thug/Analysis/screenshot/Screenshot.py index 5701d1cf65..58d393c3a9 100644 --- a/thug/Analysis/screenshot/Screenshot.py +++ b/thug/Analysis/screenshot/Screenshot.py @@ -23,7 +23,7 @@ def run(self, window, url, response, ctype): return if not ctype.startswith(self.content_types): - return + return # pragma: no cover with sync_playwright() as p: browser = p.chromium.launch(headless=True) @@ -39,8 +39,14 @@ def block_resource_type(route): try: page.set_content(response.text) + + # 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") + screenshot = page.screenshot(type="png", full_page=True) browser.close() log.ThugLogging.log_screenshot(url, screenshot) except Exception as e: - log.warning("[SCREENSHOT] Error: %s", str(e)) + log.warning("[SCREENSHOT] Error: %s", str(e)) # pragma: no cover From da851392c95d76d352428e2192fbe76c6ea6fb68 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Fri, 11 Sep 2026 11:47:56 +0200 Subject: [PATCH 7/9] Remove code from coverage --- thug/Analysis/screenshot/Screenshot.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/thug/Analysis/screenshot/Screenshot.py b/thug/Analysis/screenshot/Screenshot.py index 58d393c3a9..00b50f763f 100644 --- a/thug/Analysis/screenshot/Screenshot.py +++ b/thug/Analysis/screenshot/Screenshot.py @@ -29,7 +29,7 @@ def run(self, window, url, response, ctype): browser = p.chromium.launch(headless=True) page = browser.new_page() - def block_resource_type(route): + def block_resource_type(route): # pragma: no cover if route.request.resource_type in self.resource_types: route.continue_() else: @@ -48,5 +48,5 @@ def block_resource_type(route): screenshot = page.screenshot(type="png", full_page=True) browser.close() log.ThugLogging.log_screenshot(url, screenshot) - except Exception as e: - log.warning("[SCREENSHOT] Error: %s", str(e)) # pragma: no cover + except Exception as e: # pragma: no cover + log.warning("[SCREENSHOT] Error: %s", str(e)) From eaf5de877d967a46da51e26879e2c7468bfd7b43 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Fri, 11 Sep 2026 12:00:47 +0200 Subject: [PATCH 8/9] Enable screenshot tests on MacOS X --- tests/functional/test_screenshot.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/functional/test_screenshot.py b/tests/functional/test_screenshot.py index 1a8d462831..ebf67f0519 100644 --- a/tests/functional/test_screenshot.py +++ b/tests/functional/test_screenshot.py @@ -7,10 +7,6 @@ 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,16 +34,10 @@ 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): 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): expected = [] self.do_perform_test( From 370e5388a907b87df37fc88c3650ae247b925cc1 Mon Sep 17 00:00:00 2001 From: Angelo Dell'Aera Date: Fri, 11 Sep 2026 12:12:06 +0200 Subject: [PATCH 9/9] Minor linting change --- tests/functional/test_screenshot.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/functional/test_screenshot.py b/tests/functional/test_screenshot.py index ebf67f0519..1d36cbc347 100644 --- a/tests/functional/test_screenshot.py +++ b/tests/functional/test_screenshot.py @@ -1,7 +1,5 @@ -import os import logging -import pytest from thug.ThugAPI.ThugAPI import ThugAPI