#!/usr/bin/env python3
"""
Persistent TLS TCP client with client cert/key authentication.
Sends "I am groot" every INTERVAL seconds.

Example:
  python3 tls_echo_client.py --host af4454b46916935b6.awsglobalaccelerator.com \
                             --port 2129 \
                             --cert pinned-client.crt \
                             --key pinned-client.key \
                             --interval 10
"""

import socket
import ssl
import time
import sys
import signal
import argparse
from typing import Optional
from datetime import datetime

stop_requested = False

def log(msg: str):
    """Print log messages with timestamps."""
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{ts}] {msg}", flush=True)

def handle_sigint(signum, frame):
    global stop_requested
    log("SIGINT received, shutting down...")
    stop_requested = True

signal.signal(signal.SIGINT, handle_sigint)

def make_ssl_context(certfile: str, keyfile: str, cafile: Optional[str], insecure: bool) -> ssl.SSLContext:
    ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
    if insecure:
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
    elif cafile:
        ctx.load_verify_locations(cafile)
    ctx.load_cert_chain(certfile=certfile, keyfile=keyfile)
    return ctx

def connect_with_retry(host: str, port: int, ctx: ssl.SSLContext, timeout=5.0):
    attempt = 0
    backoff = 1.0
    while not stop_requested:
        try:
            attempt += 1
            log(f"Connecting to {host}:{port} (attempt {attempt})...")
            raw_sock = socket.create_connection((host, port), timeout=timeout)
            server_hostname = host if ctx.check_hostname else None
            ssl_sock = ctx.wrap_socket(raw_sock, server_hostname=server_hostname)
            log(f"TLS handshake complete. Cipher: {ssl_sock.cipher()}")
            return ssl_sock
        except Exception as e:
            log(f"Connection failed: {e!r}. Retrying in {backoff:.1f}s ...")
            time.sleep(backoff)
            backoff = min(backoff * 2, 30.0)
    raise KeyboardInterrupt("Stop requested")

def try_recv(ssl_sock: ssl.SSLSocket, timeout: float = 2.0) -> bytes:
    ssl_sock.settimeout(timeout)
    try:
        data = ssl_sock.recv(4096)
        return data
    except socket.timeout:
        return b""
    except Exception:
        raise

def main():
    parser = argparse.ArgumentParser(description="Persistent TLS TCP client with client cert/key")
    parser.add_argument("--host", default="af4454b46916935b6.awsglobalaccelerator.com")
    parser.add_argument("--port", type=int, default=2129)
    parser.add_argument("--cert", default="pinned-client.crt")
    parser.add_argument("--key", default="pinned-client.key")
    parser.add_argument("--cafile", default=None, help="CA file to trust in addition to the system store (optional)")
    parser.add_argument("--insecure", action="store_true", help="Skip server certificate and hostname verification")
    parser.add_argument("--interval", type=int, default=10, help="Seconds between sends")
    parser.add_argument("--message", default="I am groot", help="Message to send each interval")
    args = parser.parse_args()

    ctx = make_ssl_context(args.cert, args.key, args.cafile, args.insecure)
    ssl_sock = None

    try:
        ssl_sock = connect_with_retry(args.host, args.port, ctx)
        last_send = 0.0

        while not stop_requested:
            now = time.time()
            if now - last_send >= args.interval:
                payload = args.message
                to_send = payload + "\n"
                try:
                    ssl_sock.sendall(to_send.encode("utf-8"))
                    log(f"Sent: {payload!r}")
                except Exception as e:
                    log(f"Send failed: {e!r}. Reconnecting...")
                    try: ssl_sock.close()
                    except: pass
                    ssl_sock = connect_with_retry(args.host, args.port, ctx)
                    last_send = time.time()
                    continue

                # Try to read echo reply
                try:
                    data = try_recv(ssl_sock)
                    if data:
                        try:
                            text = data.decode("utf-8", errors="replace")
                        except Exception:
                            text = repr(data)
                        log(f"Received ({len(data)} bytes): {text!r}")
                    else:
                        log("No immediate reply.")
                except Exception as e:
                    log(f"Receive error: {e!r}. Reconnecting...")
                    try: ssl_sock.close()
                    except: pass
                    ssl_sock = connect_with_retry(args.host, args.port, ctx)

                last_send = time.time()
            time.sleep(0.2)

    except KeyboardInterrupt:
        pass
    finally:
        if ssl_sock:
            try: ssl_sock.close()
            except: pass
        log("Client exited.")

if __name__ == "__main__":
    main()