Spaces:
Sleeping
Sleeping
File size: 6,235 Bytes
748113b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | import time
from flask import Flask, request, Response, session, jsonify
from flask_socketio import SocketIO, emit
from werkzeug.serving import make_server
import requests
import re
import logging
# Configure logging
logger = logging.getLogger(__name__)
class ProxyServer:
def __init__(self, secret_key, host='localhost', port=5000):
self.host = host
self.port = port
self.app = Flask(__name__)
self.app.secret_key = secret_key
self.server = None
self.is_running = False
self.socketio = SocketIO(self.app) # Initialize SocketIO with the Flask app
# self.server_thread = None # Thread for running the server
# self.server_running = False # Flag to track server state
self.setup_routes()
self.highlight_word = None # Initialize the highlight word
@self.app.route('/shutdown', methods=['POST'])
def shutdown():
logger.info("Shutdown request received")
self.shutdown_server()
return 'Server shutting down...'
logger.info("Proxy server initialized")
# Inject JavaScript into HTML content to highlight words and listen for WebSocket updates
def inject_script(self, content):
# Inject the WebSocket listening script
script = f"""
<script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.0.1/socket.io.js"></script>
<script>
let currentHighlight = "{self.highlight_word}";
function highlightWord(word) {{
if (word) {{
document.body.innerHTML = document.body.innerHTML.replace(
new RegExp(word, 'g'),
'<span style="background-color: yellow;">' + word + '</span>'
);
}}
}}
highlightWord(currentHighlight);
// Connect to WebSocket
const socket = io();
socket.on('new_highlight', function(data) {{
currentHighlight = data.highlight;
highlightWord(currentHighlight);
}});
</script>
"""
return re.sub(r'</body>', script + '</body>', content)
# Ensure the target_url and path are handled correctly
def build_full_url(self, target_url, path):
if not target_url.endswith('/') and not path.startswith('/'):
return f"{target_url}/{path}"
return f"{target_url}{path}"
# Route handler for proxying requests
def proxy(self, path=''):
target_url = request.args.get('target_url')
if not target_url and 'target_url' in session:
target_url = session['target_url']
elif target_url:
session['target_url'] = target_url
if not target_url:
logger.error("No target_url provided")
return "Error: target_url query parameter is required", 400
full_target_url = self.build_full_url(target_url, path)
logger.info(f"Proxying request to {full_target_url}")
headers = {key: value for key, value in request.headers if key != 'Host'}
# Handle POST or GET requests
if request.method == 'POST':
response = requests.post(full_target_url, headers=headers, data=request.get_data(), stream=True)
else:
response = requests.get(full_target_url, headers=headers, stream=True)
# If it's HTML content, inject the script
if 'text/html' in response.headers.get('Content-Type', ''):
def generate():
for chunk in response.iter_content(chunk_size=1024):
if chunk:
rewritten_chunk = self.inject_script(chunk.decode('utf-8'))
yield rewritten_chunk.encode('utf-8')
logger.info(f"Injecting script into HTML response from {full_target_url}")
return Response(generate(), content_type=response.headers['Content-Type'])
# Stream non-HTML content (images, scripts, etc.)
else:
def generate():
for chunk in response.iter_content(chunk_size=1024):
if chunk:
yield chunk
return Response(generate(), content_type=response.headers['Content-Type'])
# API endpoint to set a new highlight word
def set_highlight(self):
new_highlight = request.json.get('highlight')
if new_highlight:
self.highlight_word = new_highlight
# Emit the new highlight word to all connected clients
self.socketio.emit('new_highlight', {'highlight': new_highlight})
logger.info(f"Highlight word updated to '{new_highlight}' and broadcasted to clients")
return jsonify({"message": "Highlight word updated", "highlight": new_highlight}), 200
logger.error("No highlight word provided")
return jsonify({"error": "No highlight word provided"}), 400
# Setup routes to proxy all requests and WebSocket events
def setup_routes(self):
self.app.add_url_rule('/', defaults={'path': ''}, view_func=self.proxy, methods=['GET', 'POST', 'PUT', 'DELETE', 'OPTIONS'])
self.app.add_url_rule('/<path:path>', view_func=self.proxy, methods=['GET', 'POST', 'PUT', 'DELETE', 'OPTIONS'])
self.app.add_url_rule('/set_highlight', view_func=self.set_highlight, methods=['POST'])
def run(self):
"""Runs the Werkzeug server"""
logging.info(f"Starting server on {self.host}:{self.port}")
self.server = make_server(self.host, self.port, self.app, threaded=True)
self.is_running = True
self.server.serve_forever()
def shutdown_server(self):
"""Shuts down the Werkzeug server"""
if self.server:
logger.info("Shutting down server...")
self.is_running = False
self.server.shutdown()
logger.info("Server shut down complete")
# Create an instance of ProxyServer and expose the app
# proxy_server = ProxyServer(secret_key='your_secret_key_here')
# app = proxy_server.app # Expose the Flask app to the top-level scope for Flask CLI
# if __name__ == '__main__':
# proxy_server.run(port=5000)
|