Update pdf2zh/gui.py

pdf2zh/gui.py

@@ -48,25 +48,22 @@ if os.environ.get("PDF2ZH_DEMO"):
         "First": [0],
         "First 20 pages": list(range(0, 20)),
     }
+    client_key = os.environ.get("PDF2ZH_CLIENT_KEY")
+    server_key = os.environ.get("PDF2ZH_SERVER_KEY")
+
+
+def verify_recaptcha(response):
+    recaptcha_url = "https://www.google.com/recaptcha/api/siteverify"
+
+    print("reCAPTCHA", server_key, response)
+
+    data = {"secret": server_key, "response": response}
+    result = requests.post(recaptcha_url, data=data).json()
+
+    print("reCAPTCHA", result.get("success"))
+
+    return result.get("success")
 
-def download_with_limit(url, save_path, size_limit):
-    chunk_size = 1024
-    total_size = 0
-    with requests.get(url, stream=True, timeout=10) as response:
-        response.raise_for_status()
-        content = response.headers.get("Content-Disposition")
-        try:
-            _, params = cgi.parse_header(content)
-            filename = params["filename"]
-        except Exception:
-            filename = os.path.basename(url)
-        with open(save_path / filename, "wb") as file:
-            for chunk in response.iter_content(chunk_size=chunk_size):
-                total_size += len(chunk)
-                if size_limit and total_size > size_limit:
-                    raise gr.Error("Exceeds file size limit")
-                file.write(chunk)
-    return save_path / filename
 
 def pdf_preview(file):
     doc = pymupdf.open(file)
@@ -75,6 +72,7 @@ def pdf_preview(file):
     image = np.frombuffer(pix.samples, np.uint8).reshape(pix.height, pix.width, 3)
     return image
 
+
 def upload_file(file, service, progress=gr.Progress()):
     """Handle file upload, validation, and initial preview."""
     if not file or not os.path.exists(file):
@@ -89,6 +87,27 @@ def upload_file(file, service, progress=gr.Progress()):
         print(f"Error converting PDF: {e}")
         return None, None
 
+
+def download_with_limit(url, save_path, size_limit):
+    chunk_size = 1024
+    total_size = 0
+    with requests.get(url, stream=True, timeout=10) as response:
+        response.raise_for_status()
+        content = response.headers.get("Content-Disposition")
+        try:
+            _, params = cgi.parse_header(content)
+            filename = params["filename"]
+        except Exception:
+            filename = os.path.basename(url)
+        with open(save_path / filename, "wb") as file:
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                total_size += len(chunk)
+                if size_limit and total_size > size_limit:
+                    raise gr.Error("Exceeds file size limit")
+                file.write(chunk)
+    return save_path / filename
+
+
 def translate(
     file_type,
     file_input,
@@ -99,9 +118,13 @@ def translate(
     lang_from,
     lang_to,
     page_range,
+    recaptcha_response,
     progress=gr.Progress(),
 ):
     """Translate PDF content using selected service."""
+    if flag_demo and not verify_recaptcha(recaptcha_response):
+        raise gr.Error("reCAPTCHA fail")
+
     progress(0, desc="Starting translation...")
 
     output = Path("pdf2zh_files")
@@ -173,6 +196,7 @@ def translate(
         gr.update(visible=True),
     )
 
+
 # Global setup
 custom_blue = gr.themes.Color(
     c50="#E8F3FF",
@@ -198,9 +222,11 @@ with gr.Blocks(
     footer {visibility: hidden}
     .env-warning {color: #dd5500 !important;}
     .env-success {color: #559900 !important;}
+    /* Add dashed border to input-file class */
     .input-file {
         border: 1.2px dashed #165DFF !important;
         border-radius: 6px !important;
+        # background-color: #ffffff !important;
         transition: background-color 0.4s ease-out;
     }
     .input-file:hover {
@@ -216,9 +242,36 @@ with gr.Blocks(
     .progress-bar {
     border-radius: 8px !important;
     }
+    # .input-file label {
+    #     color: #165DFF !important;
+    #     border: 1.2px dashed #165DFF !important;
+    #     border-left: none !important;
+    #     border-top: none !important;
+    # }
+    # .input-file .wrap {
+    #     color: #165DFF !important;
+    # }
+    # .input-file .or {
+    #     color: #165DFF !important;
+    # }
     """,
+    head=(
+        """
+    <script src="https://www.google.com/recaptcha/api.js?render=explicit" async defer></script>
+    <script type="text/javascript">
+        var onVerify = function(token) {
+            el=document.getElementById('verify').getElementsByTagName('textarea')[0];
+            el.value=token;
+            el.dispatchEvent(new Event('input'));
+        };
+    </script>
+    """
+        if flag_demo
+        else ""
+    ),
 ) as demo:
     gr.Markdown(
+        #"# [PDFMathTranslate @ GitHub](https://github.com/Byaidu/PDFMathTranslate)"
         "# [PDFMathTranslate——科研之心免费提供（更多科研AI智能体请点击）](https://ai.linkagi.top)"
     )
 
@@ -337,6 +390,10 @@ with gr.Blocks(
             output_file_dual = gr.File(
                 label="Download Translation (Dual)", visible=False
             )
+            recaptcha_response = gr.Textbox(
+                label="reCAPTCHA Response", elem_id="verify", visible=False
+            )
+            recaptcha_box = gr.HTML('<div id="recaptcha-box"></div>')
             translate_btn = gr.Button("Translate", variant="primary")
             tech_details_tog = gr.Markdown(
                 details_wrapper(envs_status),
@@ -349,6 +406,21 @@ with gr.Blocks(
                 on_select_filetype,
                 file_type,
                 [file_input, link_input],
+                js=(
+                    f"""
+                    (a,b)=>{{
+                        try{{
+                            grecaptcha.render('recaptcha-box',{{
+                                'sitekey':'{client_key}',
+                                'callback':'onVerify'
+                            }});
+                        }}catch(error){{}}
+                        return [a];
+                    }}
+                    """
+                    if flag_demo
+                    else ""
+                ),
             )
 
         with gr.Column(scale=2):
@@ -360,6 +432,21 @@ with gr.Blocks(
         upload_file,
         inputs=[file_input, service],
         outputs=[file_input, preview],
+        js=(
+            f"""
+            (a,b)=>{{
+                try{{
+                    grecaptcha.render('recaptcha-box',{{
+                        'sitekey':'{client_key}',
+                        'callback':'onVerify'
+                    }});
+                }}catch(error){{}}
+                return [a];
+            }}
+            """
+            if flag_demo
+            else ""
+        ),
     )
 
     translate_btn.click(
@@ -374,6 +461,7 @@ with gr.Blocks(
             lang_from,
             lang_to,
             page_range,
+            recaptcha_response,
         ],
         outputs=[
             output_file,
@@ -383,7 +471,8 @@ with gr.Blocks(
             output_file_dual,
             output_title,
         ],
-    )
+    ).then(lambda: None, js="()=>{grecaptcha.reset()}" if flag_demo else "")
+
 
 def setup_gui(share=False):
     if flag_demo:
@@ -405,6 +494,7 @@ def setup_gui(share=False):
                 )
                 demo.launch(debug=True, inbrowser=True, share=True)
 
+
 # For auto-reloading while developing
 if __name__ == "__main__":
     setup_gui()