proxy_server.py

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)