Update app.py

app.py

@@ -25,6 +25,10 @@ import shutil
 from PIL import Image
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
+import google_auth_oauthlib.flow
+import googleapiclient.discovery
+import google.auth.transport.requests
+import googleapiclient.http
 
 # -------------------- Logging Setup --------------------
 logging.basicConfig(
@@ -33,6 +37,19 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# -------------------- Google OAuth Config --------------------
+GOOGLE_OAUTH_CONFIG = {
+    "web": {
+        "client_id": "90798824947-u25obg1q844qeikjoh4jdmi579kn9p1c.apps.googleusercontent.com",
+        "project_id": "huggingface-449214",
+        "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+        "token_uri": "https://oauth2.googleapis.com/token",
+        "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+        "client_secret": "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",
+        "redirect_uris": ["https://euler314-craw-web.hf.space/"]
+    }
+}
+
 # -------------------- Utility Functions --------------------
 def get_random_user_agent():
     USER_AGENTS = [
@@ -58,6 +75,56 @@ def create_zip_file(file_paths, output_dir):
             zipf.write(file_path, os.path.basename(file_path))
     return zip_path
 
+# -------------------- Google Drive Functions --------------------
+def get_google_auth_url():
+    client_config = GOOGLE_OAUTH_CONFIG["web"]
+    flow = google_auth_oauthlib.flow.Flow.from_client_config(
+        {"web": client_config},
+        scopes=["https://www.googleapis.com/auth/drive.file"]
+    )
+    flow.redirect_uri = client_config["redirect_uris"][0]
+    authorization_url, _ = flow.authorization_url(
+        access_type="offline",
+        include_granted_scopes="true",
+        prompt="consent"
+    )
+    return authorization_url
+
+def exchange_code_for_credentials(auth_code):
+    if not auth_code.strip():
+        return None, "No code provided."
+    try:
+        client_config = GOOGLE_OAUTH_CONFIG["web"]
+        flow = google_auth_oauthlib.flow.Flow.from_client_config(
+            {"web": client_config},
+            scopes=["https://www.googleapis.com/auth/drive.file"]
+        )
+        flow.redirect_uri = client_config["redirect_uris"][0]
+        flow.fetch_token(code=auth_code.strip())
+        creds = flow.credentials
+        if not creds or not creds.valid:
+            return None, "Could not validate credentials. Check code and try again."
+        return creds, "Google Sign-In successful!"
+    except Exception as e:
+        return None, f"Error during token exchange: {e}"
+
+def google_drive_upload(file_path, credentials, folder_id=None):
+    try:
+        drive_service = googleapiclient.discovery.build("drive", "v3", credentials=credentials)
+        file_metadata = {'name': os.path.basename(file_path)}
+        if folder_id:
+            file_metadata['parents'] = [folder_id]
+        media = googleapiclient.http.MediaFileUpload(file_path, resumable=True)
+        created = drive_service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+        return created.get("id", "")
+    except Exception as e:
+        return f"Error uploading to Drive: {str(e)}"
+
+def create_drive_folder(drive_service, name):
+    folder_metadata = {'name': name, 'mimeType': 'application/vnd.google-apps.folder'}
+    folder = drive_service.files().create(body=folder_metadata, fields='id').execute()
+    return folder.get('id')
+
 # -------------------- Playwright Setup --------------------
 def install_playwright_dependencies():
     try:
@@ -65,9 +132,17 @@ def install_playwright_dependencies():
         os.environ['PLAYWRIGHT_BROWSERS_PATH'] = os.path.expanduser("~/.cache/ms-playwright")
         
         # Install system dependencies
+        subprocess.run(['apt-get', 'update', '-y'], check=True)
+        packages = [
+            'libnss3', 'libnss3-tools', 'libnspr4', 'libatk1.0-0',
+            'libatk-bridge2.0-0', 'libatspi2.0-0', 'libcups2', 'libxcomposite1',
+            'libxdamage1', 'libdrm2', 'libgbm1', 'libpango-1.0-0'
+        ]
+        subprocess.run(['apt-get', 'install', '-y', '--no-install-recommends'] + packages, check=True)
+        
+        # Install Playwright and dependencies
         subprocess.run(['pip', 'install', 'playwright'], check=True)
-        subprocess.run(['playwright', 'install', 'chromium'], check=True)
-        subprocess.run(['playwright', 'install-deps', 'chromium'], check=True)
+        subprocess.run(['python3', '-m', 'playwright', 'install', 'chromium'], check=True)
         
         st.success("Playwright dependencies installed successfully!")
     except Exception as e:
@@ -1664,6 +1739,7 @@ def main():
         st.session_state.initialized = True
         st.session_state.discovered_files = []
         st.session_state.current_url = None
+        st.session_state.google_creds = None
         st.session_state.selected_files = []
         st.session_state.do_deep_search = False
         st.session_state.deep_search_url = None
@@ -1678,6 +1754,16 @@ def main():
             use_proxy = st.checkbox("Use Proxy", key="proxy_checkbox")
             proxy = st.text_input("Proxy URL", placeholder="http://proxy:port", key="proxy_input")
 
+        with st.expander("Google Drive Integration", expanded=False):
+            if st.button("Start Google Sign-In", key="google_signin_btn"):
+                auth_url = get_google_auth_url()
+                st.markdown(f"[Click here to authorize]({auth_url})")
+            auth_code = st.text_input("Enter authorization code", key="auth_code_input")
+            if st.button("Complete Sign-In", key="complete_signin_btn") and auth_code:
+                creds, msg = exchange_code_for_credentials(auth_code)
+                st.session_state.google_creds = creds
+                st.write(msg)
+
     if mode == "Manual URL":
         st.header("Manual URL Mode")
         url = st.text_input("Enter URL", placeholder="https://example.com", key="url_input")
@@ -1744,8 +1830,15 @@ def main():
             st.session_state.selected_files = selected_indices
             
             if selected_indices:
-                download_dir = st.text_input("Download Directory", value="./downloads", key="download_dir_input")
-                create_zip = st.checkbox("Create ZIP file", value=True, key="create_zip_checkbox")
+                col1, col2, col3, col4 = st.columns(4)
+                with col1:
+                    download_dir = st.text_input("Download Directory", value="./downloads", key="download_dir_input")
+                with col2:
+                    create_zip = st.checkbox("Create ZIP file", value=True, key="create_zip_checkbox")
+                with col3:
+                    delete_after = st.checkbox("Delete after creating ZIP", key="delete_after_checkbox")
+                with col4:
+                    upload_to_drive = st.checkbox("Upload to Google Drive", key="upload_drive_checkbox")
                 
                 if st.button("Download Selected", key="download_btn"):
                     if not os.path.exists(download_dir):
@@ -1777,8 +1870,7 @@ def main():
                     if downloaded:
                         st.success(f"Successfully downloaded {len(downloaded)} files")
                         
-                        # Create file downloads
-                        if create_zip and len(downloaded) > 1:
+                        if create_zip:
                             zip_path = create_zip_file(downloaded, download_dir)
                             st.success(f"Created ZIP file: {zip_path}")
                             
@@ -1786,14 +1878,32 @@ def main():
                             with open(zip_path, "rb") as f:
                                 zip_data = f.read()
                             
-                            zip_filename = os.path.basename(zip_path)
                             st.download_button(
                                 label="Download ZIP",
                                 data=zip_data,
-                                file_name=zip_filename,
+                                file_name=os.path.basename(zip_path),
                                 mime="application/zip",
                                 key="download_zip_btn"
                             )
+                            
+                            # Upload to Google Drive if requested
+                            if upload_to_drive and st.session_state.google_creds:
+                                drive_service = googleapiclient.discovery.build("drive", "v3", credentials=st.session_state.google_creds)
+                                folder_id = create_drive_folder(drive_service, f"Downloads_{urlparse(url).netloc}")
+                                drive_id = google_drive_upload(zip_path, st.session_state.google_creds, folder_id)
+                                if not isinstance(drive_id, str) or not drive_id.startswith("Error"):
+                                    st.success(f"Uploaded to Google Drive. File ID: {drive_id}")
+                                else:
+                                    st.error(drive_id)
+                                    
+                            # Delete original files if requested
+                            if delete_after:
+                                for path in downloaded:
+                                    try:
+                                        os.remove(path)
+                                    except Exception as e:
+                                        st.warning(f"Could not delete {path}: {e}")
+                                st.info("Deleted original files after ZIP creation")
                         else:
                             # Provide individual file downloads
                             st.write("Download files individually:")