From 7ff4e4f44f8c3c94d0198617e4d39fd786b8b67b Mon Sep 17 00:00:00 2001
From: estebanthi <esteban.thilliez@gmail.com>
Date: Fri, 24 Jan 2025 11:08:41 +0100
Subject: [PATCH] Screenshot parameter

---
 src/dtos.py                 | 2 ++
 src/flaresolverr_service.py | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/src/dtos.py b/src/dtos.py
index 1e9aace..bb38279 100644
--- a/src/dtos.py
+++ b/src/dtos.py
@@ -10,6 +10,7 @@ class ChallengeResolutionResultT:
     response: str = None
     cookies: list = None
     userAgent: str = None
+    screenshot: str | None = None
 
     def __init__(self, _dict):
         self.__dict__.update(_dict)
@@ -41,6 +42,7 @@ class V1RequestBase(object):
     url: str = None
     postData: str = None
     returnOnlyCookies: bool = None
+    returnScreenshot: bool = None
     download: bool = None   # deprecated v2.0.0, not used
     returnRawHtml: bool = None  # deprecated v2.0.0, not used
 
diff --git a/src/flaresolverr_service.py b/src/flaresolverr_service.py
index cfc2088..e1d4152 100644
--- a/src/flaresolverr_service.py
+++ b/src/flaresolverr_service.py
@@ -418,6 +418,9 @@ def _evil_logic(req: V1RequestBase, driver: WebDriver, method: str) -> Challenge
         challenge_res.headers = {}  # todo: fix, selenium not provides this info
         challenge_res.response = driver.page_source
 
+    if req.returnScreenshot:
+        challenge_res.screenshot = driver.get_screenshot_as_base64()
+
     res.result = challenge_res
     return res