#!/usr/bin/env python3
import http.server
import socket
import socketserver
import ipaddress
import logging
import argparse

# Configure logging (thread-safe), output time and level to stdout
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

class TCPServer6(socketserver.TCPServer):
    address_family = socket.AF_INET6
    allow_reuse_address = True

class MyHTTPRequestHandler(http.server.SimpleHTTPRequestHandler):
    def log_response(self, status_code, client_ip, x_forwarded_for):
        logging.info(
            "method=%s path=%s client_ip=%s x_forwarded_for=%s status=%d",
            self.command,
            self.path,
            client_ip,
            x_forwarded_for if x_forwarded_for is not None else "-",
            status_code,
        )

    def do_GET(self):
        # Get 'X-Forwarded-For' header or remote address
        x_forwarded_for = self.headers.get('X-Forwarded-For')
        remote_addr = self.client_address[0]

        # Prefer client IP from X-Forwarded-For if present
        client_ip = (x_forwarded_for.split(',')[0].strip() if x_forwarded_for else remote_addr)

        # Convert IPv4-mapped IPv6 addresses to plain IPv4 (e.g. ::ffff:127.0.0.1 -> 127.0.0.1)
        try:
            ip_obj = ipaddress.ip_address(client_ip)
            if isinstance(ip_obj, ipaddress.IPv6Address) and ip_obj.ipv4_mapped:
                client_ip = str(ip_obj.ipv4_mapped)
        except ValueError:
            # If parsing fails (non-standard format), attempt manual handling to cope with some proxies
            # Remove IPv6 brackets and port (e.g. [::ffff:127.0.0.1]:8080)
            s = client_ip
            if s.startswith('[') and ']' in s:
                s = s.split(']', 1)[0].lstrip('[')
            if s.count(':') >= 2 and s.startswith('::ffff:'):
                # Handle possible port by taking the last segment
                s = s.split(':')[-1]
                client_ip = s

        # Send response
        status_code = 200
        self.send_response(status_code)
        self.send_header("Content-type", "text/plain")
        self.end_headers()
        self.wfile.write(client_ip.encode())

        # Log after responding
        try:
            self.log_response(status_code, client_ip, x_forwarded_for)
        except Exception:
            logging.exception("Failed to log response")

def parse_args():
    parser = argparse.ArgumentParser(description="Simple IPv6-capable HTTP server that returns client IP.")
    parser.add_argument("-p", "--port", type=int, default=8080, help="Port to listen on (default: 8080)")
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    PORT = args.port

    # For AF_INET6 bind to all IPv6 addresses using a 4-tuple
    server_address = ("::", PORT, 0, 0)
    with TCPServer6(server_address, MyHTTPRequestHandler) as httpd:
        # Optionally set IPV6_V6ONLY. Set to 0 to accept IPv4-mapped addresses as well.
        try:
            httpd.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
        except Exception:
            pass
        logging.info("Serving on [::]:%d", PORT)
        httpd.serve_forever()