How to Intercept and Forward Cookies in Python Flask

Question

I'm working on a setup where a Flask application (main server) is accessed through a proxy server. The goal is to leverage the main server's session management functionality for user interactions.

Problem:

The current behavior seems to result in a new session being created on the main server for each request that goes through the proxy. This disrupts session continuity and prevents features like tracking user actions or maintaining login state.

Desired Outcome:

I want the proxy to intercept the initial Flask session cookie set by the main server in the first request. Subsequently, the proxy should include this cookie in all future requests to the main server, enabling session persistence and proper session handling.

Question:

How can I effectively intercept the initial Flask session cookie set by the main server on the first request through the proxy? What's the best approach to store and forward the session cookie in the proxy for subsequent requests to maintain the session on the main server? Verification:

To confirm successful session persistence, I'll be looking for the following:

Consistent Internal Tracking ID: The response from the main server should return the same internal tracking ID on subsequent requests within the same session. Session Counter Increment: The session-related counter maintained on the main server should increment appropriately by 1 with each request within the same session.

Below is the code I have been using:

import logging
import requests
from urllib.parse import urlparse, urlunparse
import uuid
from flask_cors import CORS

app = Flask(__name__)
CORS(app)  # Initialize CORS for the entire application
logging.basicConfig(level=logging.INFO)

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
APPROVED_HOSTS = {"google.com", "www.google.com", "yahoo.com", "www.trackstage.io"}
CHUNK_SIZE = 1024

# Session handling utilities
SESSION_STORE = {}  # Dictionary to store session IDs

def get_session_id(client_ip):
    """
    Get a session ID for the given client IP address.
    If a session ID exists for the client, return it.
    Otherwise, generate a new session ID and store it.
    """
    if client_ip in SESSION_STORE:
        session_id = SESSION_STORE[client_ip]
    else:
        session_id = str(uuid.uuid4())
        SESSION_STORE[client_ip] = session_id
        logging.info("New session created for client %s: %s", client_ip, session_id)
    return session_id

@app.before_request
def before_request():
    # Get the client's IP address
    client_ip = request.remote_addr

    # Get or create a session ID for the client
    session_id = get_session_id(client_ip)

    # Set the session ID in the session object
    session['session_id'] = session_id

@app.route('/p/<path:url>')
def proxy(url):
    # Prepend the scheme if missing
    if not urlparse(url).scheme:
        url = 'http://' + url

    # Get the session ID from the session object
    session_id = session.get('session_id')

    # Set session cookie in outgoing request
    headers = {'Cookie': f'session={session_id}'}

    # Send the request and get the response
    response = requests.get(url, headers=headers)

    # Update the response's Set-Cookie header
    set_cookie_header = response.headers.get('Set-Cookie')
    if set_cookie_header:
        response.headers['Set-Cookie'] = f'session={session_id}; {set_cookie_header.split(";", 1)[1]}'

    # Return the response
    return response.content


def get_source_rsp(url):
    if not urlparse(url).scheme:
        url = 'http://' + url
    logging.info("Fetching %s", url)
    if not is_approved(url):
        logging.warning("URL is not approved: %s", url)
        abort(403)
    parsed_url = urlparse(url)
    url_without_params = urlunparse((parsed_url.scheme, parsed_url.netloc, parsed_url.path, '', '', ''))

    # Get session cookie
    session_cookie = SessionCookieHandler.get_session_cookie()
    headers = {}
    if session_cookie:  # If session cookie is already stored for the client
        # Use the stored session cookie for subsequent hits
        headers['Cookie'] = session_cookie

    # Copy important headers from original request
    headers['User-Agent'] = request.headers.get('User-Agent')
    headers['Accept'] = request.headers.get('Accept')
    headers['Sec-Fetch-Dest'] = request.headers.get('Sec-Fetch-Dest')
    headers['Sec-Fetch-Mode'] = request.headers.get('Sec-Fetch-Mode')
    headers['Sec-Fetch-Site'] = request.headers.get('Sec-Fetch-Site')
    headers['Sec-Fetch-User'] = request.headers.get('Sec-Fetch-User')

    logging.info("Request headers for client %s: %s", request.remote_addr, headers)

    # Send the request and get the response
    response = requests.get(url_without_params, stream=True, headers=headers)
    
    # Update session cookie from response
    if 'Set-Cookie' in response.headers:
        set_cookie_header = response.headers['Set-Cookie']
        session_cookie = parse_session_cookie(set_cookie_header)
        if session_cookie:
            SessionCookieHandler.set_session_cookie(session_cookie)
            logging.info("Updated session cookie from response: %s", session_cookie)

    response.raise_for_status()  # Raise an error for non-2xx responses
    return response


def parse_session_cookie(set_cookie_header):
    """Parse session cookie from Set-Cookie header."""
    cookies = set_cookie_header.split(';')
    for cookie in cookies:
        parts = cookie.split('=')
        if len(parts) == 2 and parts[0].strip() == 'session':
            return parts[1].strip()
    return None


def is_approved(url):
    """Indicates whether the given URL is allowed to be fetched."""
    host = split_url(url)[1]
    return host in APPROVED_HOSTS

def split_url(url):
    """Splits the given URL into a tuple of (protocol, host, uri)."""
    proto, rest = url.split(':', 1)
    rest = rest[2:].split('/', 1)
    host, uri = (rest[0], rest[1]) if len(rest) == 2 else (rest[0], "")
    return proto, host, uri

def stream_response(response):
    """Streams the response content."""
    return response.text

if __name__ == "__main__":
    app.secret_key = 'super secret key'  # Set the secret key to use sessions
    app.run(debug=True) ```