pqs_chat.py

import argparse, socket, struct, threading, time, os, ipaddress, sys
import pqs_core as PQS

BANNER = b"Petoron P2P Messenger / Ivan Alekseev"

FRAME_LEN_SZ       = 4
HANDSHAKE_TIMEOUT  = 8.0
IDLE_TIMEOUT       = 0.0
ACK_TIMEOUT        = 5.0
SLOWLORIS_READ_TMO = 1.5
MAX_FRAME          = getattr(PQS, "MAX_DATA_SIZE", 10**7)
REPLAY_CACHE_MAX   = 4000

T_MSG  = 0x01
T_ACK  = 0x02
T_NOISE= 0x10

def now_ts() -> int:
    return int(time.time())

def pack_frame(b: bytes) -> bytes:
    return struct.pack(">I", len(b)) + b

def recv_exact(s: socket.socket, n: int) -> bytes:
    buf = bytearray()
    deadline = None
    while len(buf) < n:
        rem = n - len(buf)

        if SLOWLORIS_READ_TMO and SLOWLORIS_READ_TMO > 0:
            if len(buf) == 0:
                try:
                    s.settimeout(None)
                except Exception:
                    pass
            else:
                if deadline is None:
                    deadline = time.time() + SLOWLORIS_READ_TMO
                left = max(0.05, deadline - time.time())
                s.settimeout(left)
        else:
            try:
                s.settimeout(None)
            except Exception:
                pass

        try:
            chunk = s.recv(rem)
        except socket.timeout:
            if len(buf) == 0:
                continue
            raise

        if not chunk:
            raise ConnectionError("peer closed")

        buf.extend(chunk)
        if SLOWLORIS_READ_TMO and SLOWLORIS_READ_TMO > 0:
            deadline = time.time() + SLOWLORIS_READ_TMO

    try:
        s.settimeout(None)
    except Exception:
        pass
    return bytes(buf)

def recv_frame(s: socket.socket) -> bytes:
    hdr = recv_exact(s, FRAME_LEN_SZ)
    (ln,) = struct.unpack(">I", hdr)
    if ln <= 0 or ln > MAX_FRAME:
        raise ValueError("bad frame size")
    return recv_exact(s, ln)

def pqs_wrap_encrypt(payload: bytes, pw: str, double: bool) -> bytes:
    ct = PQS.pqs_encrypt(payload, pw)
    if double:
        ct = PQS.pqs_encrypt(ct, pw)
    return ct

def pqs_wrap_decrypt(blob: bytes, pw: str, double: bool) -> bytes:
    if not double:
        return PQS.pqs_decrypt(blob, pw)
    x = PQS.pqs_decrypt(blob, pw)
    return PQS.pqs_decrypt(x, pw)

def pqs_send(s: socket.socket, pw: str, payload: bytes, stealth: bool):
    s.sendall(pack_frame(pqs_wrap_encrypt(payload, pw, stealth)))

def pqs_recv(s: socket.socket, pw: str, stealth: bool) -> bytes:
    return pqs_wrap_decrypt(recv_frame(s), pw, stealth)

def is_bind_ok(ip: str) -> bool:
    if ip in ("0.0.0.0", "::"): return True
    try:
        a = ipaddress.ip_address(ip)
        return a.is_global and not (a.is_loopback or a.is_link_local or a.is_multicast or a.is_unspecified or a.is_reserved or a.is_private)
    except ValueError:
        return False

def require_connect_ip(ip: str) -> str:
    a = ipaddress.ip_address(ip)
    if not a.is_global or a.is_private or a.is_loopback or a.is_link_local or a.is_multicast or a.is_unspecified or a.is_reserved:
        raise ValueError("global IP required")
    return ip

def set_keepalive(sock: socket.socket):
    try:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
        if hasattr(socket, "TCP_KEEPIDLE"):  sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30)
        if hasattr(socket, "TCP_KEEPINTVL"): sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 10)
        if hasattr(socket, "TCP_KEEPCNT"):   sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 3)
    except Exception:
        pass

def do_handshake_server(c: socket.socket, pw: str, stealth: bool) -> None:
    c.settimeout(HANDSHAKE_TIMEOUT)
    hello = pqs_recv(c, pw, stealth=False)
    if not hello.startswith(b"HELLO:"):
        raise ValueError("bad hello")
    cn = hello.split(b":", 1)[1]
    sn = os.urandom(12)
    pqs_send(c, pw, b"WELCOME:" + cn + b":" + sn, stealth=False)
    ok = pqs_recv(c, pw, stealth=False)
    if ok != b"OK:" + sn:
        raise ValueError("bad OK")
    pqs_send(c, pw, BANNER, stealth=False)
    c.settimeout(None)

def do_handshake_client(c: socket.socket, pw: str, stealth: bool) -> None:
    c.settimeout(HANDSHAKE_TIMEOUT)
    cn = os.urandom(12)
    pqs_send(c, pw, b"HELLO:" + cn, stealth=False)
    w = pqs_recv(c, pw, stealth=False)
    if not w.startswith(b"WELCOME:"):
        raise ValueError("bad welcome")
    _, rest = w.split(b":", 1)
    gotc, sn = rest.split(b":", 1)
    if gotc != cn:
        raise ValueError("nonce mismatch")
    pqs_send(c, pw, b"OK:" + sn, stealth=False)
    bnr = pqs_recv(c, pw, stealth=False)
    if bnr != BANNER:
        raise ValueError("banner mismatch")
    c.settimeout(None)

def build_msg(nick_b: bytes, body_b: bytes, pad_min: int, pad_max: int):
    mid = os.urandom(16)
    ts  = now_ts()
    nlen = len(nick_b)
    blen = len(body_b)
    pad_len = 0
    pad = b""
    if pad_max > 0 and pad_max >= pad_min:
        pad_len = int.from_bytes(os.urandom(2), "big") % (pad_max - pad_min + 1) + pad_min
        pad = os.urandom(pad_len)
    hdr = bytes([T_MSG]) + mid + struct.pack(">QHI", ts, nlen, blen)
    pkt = hdr + nick_b + body_b + struct.pack(">H", pad_len) + pad
    return mid, pkt

def parse_msg(blob: bytes):
    if len(blob) < 1: return None
    if blob[0] != T_MSG: return None
    off = 1
    if len(blob) < off+16+8+2+4: return None
    mid = blob[off:off+16]; off += 16
    ts, nlen, blen = struct.unpack(">QHI", blob[off:off+14]); off += 14
    need = off + nlen + blen + 2
    if len(blob) < need: return None
    nick_b = blob[off:off+nlen]; off += nlen
    body_b = blob[off:off+blen]; off += blen
    pad_len = struct.unpack(">H", blob[off:off+2])[0]; off += 2
    if len(blob) != off + pad_len: return None
    return mid, ts, nick_b, body_b

def build_ack(mid: bytes) -> bytes:
    return bytes([T_ACK]) + mid

def parse_ack(blob: bytes):
    if len(blob) != 1+16: return None
    if blob[0] != T_ACK: return None
    return blob[1:]

def build_noise():
    n = 8 + (os.urandom(1)[0] % 56)
    return bytes([T_NOISE]) + os.urandom(n)

def is_noise(blob: bytes) -> bool:
    return len(blob) >= 1 and blob[0] == T_NOISE

class Chat:
    def __init__(self, sock: socket.socket, peer: str, pw: str, nick_bytes: bytes, stealth: bool):
        self.s, self.peer, self.password, self.stealth = sock, peer, pw, stealth
        self.nick = nick_bytes
        self.alive = True
        self.pending = {}
        self.plock   = threading.Lock()
        self.seen_ids= set()
        if self.stealth:
            self.pad_min, self.pad_max = 8, 64
        else:
            self.pad_min, self.pad_max = 0, 0

    def mark_sent(self, mid: bytes):
        with self.plock:
            self.pending[mid] = {"ts": now_ts()}
        sys.stdout.write("→ [Sent] ")
        sys.stdout.flush()

    def mark_ack(self, mid: bytes):
        with self.plock:
            it = self.pending.pop(mid, None)
        if it is not None:
            sys.stdout.write("→ [status] Delivered\n")
            sys.stdout.flush()

    def mark_fail(self, mid: bytes):
        with self.plock:
            it = self.pending.pop(mid, None)
        if it is not None:
            sys.stdout.write("→ [status] Not delivered\n")
            sys.stdout.flush()

    def ack_watcher(self):
        try:
            while self.alive:
                time.sleep(0.2)
                nowv = now_ts()
                drops = []
                with self.plock:
                    for mid, it in list(self.pending.items()):
                        if nowv - it["ts"] > ACK_TIMEOUT:
                            drops.append(mid)
                for mid in drops:
                    self.mark_fail(mid)
        except Exception:
            pass

    def noise_pumper(self):
        if not self.stealth:
            return
        try:
            while self.alive:
                time.sleep(1.5 + (os.urandom(1)[0] % 250) / 100.0)
                try:
                    pqs_send(self.s, self.password, build_noise(), stealth=True)
                except Exception:
                    pass
        except Exception:
            pass

    def sender(self):
        try:
            while self.alive:
                try:
                    line = sys.stdin.buffer.readline()
                except Exception as e:
                    sys.stdout.write(f"sender exception: {e}\n")
                    sys.stdout.flush()
                    time.sleep(0.2)
                    continue

                if not line:
                    time.sleep(0.05)
                    continue

                line = line.rstrip(b"\r\n")
                if line == b"/quit":
                    self.stop()
                    break

                mid, pkt = build_msg(self.nick, line, self.pad_min, self.pad_max)
                try:
                    if not self.alive:
                        break
                    if getattr(self.s, "fileno", lambda: -1)() < 0:
                        raise OSError("socket closed")
                    pqs_send(self.s, self.password, pkt, stealth=self.stealth)
                    self.mark_sent(mid)
                except Exception as e:
                    sys.stdout.write(f"send error: {e}\n")
                    sys.stdout.write("→ [status] Error\n")
                    sys.stdout.flush()
        except Exception as e:
            sys.stdout.write(f"sender outer exception: {e}\n")
            sys.stdout.flush()

    def receiver(self):
        try:
            while self.alive:
                try:
                    d = pqs_recv(self.s, self.password, stealth=self.stealth)
                except Exception as e:
                    sys.stdout.write(f"recv exception: {e}\n")
                    sys.stdout.flush()
                    break

                try:
                    if is_noise(d):
                        continue

                    a = parse_ack(d)
                    if a is not None:
                        self.mark_ack(a)
                        continue

                    m = parse_msg(d)
                    if m is not None:
                        mid, ts, nick_b, body_b = m
                        if mid in self.seen_ids:
                            continue
                        if len(self.seen_ids) > REPLAY_CACHE_MAX:
                            self.seen_ids.clear()
                        self.seen_ids.add(mid)

                        sys.stdout.buffer.write(b"[" + str(ts).encode("ascii") + b"] " + nick_b + b": " + body_b + b"\n")
                        sys.stdout.flush()

                        try:
                            pqs_send(self.s, self.password, build_ack(mid), stealth=self.stealth)
                        except Exception:
                            pass
                        continue

                    sys.stdout.write(f"(recv {len(d)} bytes, skipped)\n")
                    sys.stdout.flush()

                except Exception as e:
                    sys.stdout.write(f"decode exception: {e}\n")
                    sys.stdout.flush()
        except Exception as e:
            sys.stdout.write(f"receiver outer exception: {e}\n")
            sys.stdout.flush()
        finally:
            self.stop()

    def stop(self):
        if not self.alive:
            return
        self.alive = False
        try: self.s.shutdown(socket.SHUT_RDWR)
        except: pass
        try: self.s.close()
        except: pass
        with self.plock:
            for mid in list(self.pending.keys()):
                self.mark_fail(mid)
            self.pending.clear()
        try:
            self.password = "\x00" * len(self.password)
        except Exception:
            pass
        sys.stdout.write("connection closed & memory wiped\n")
        sys.stdout.flush()

    def run(self):
        sys.stdout.write(f"connected to {self.peer}\n")
        sys.stdout.flush()
        t1 = threading.Thread(target=self.sender, daemon=True)
        t2 = threading.Thread(target=self.receiver, daemon=True)
        t3 = threading.Thread(target=self.ack_watcher, daemon=True)
        t4 = threading.Thread(target=self.noise_pumper, daemon=True)
        t1.start(); t2.start(); t3.start(); t4.start()
        while any(t.is_alive() for t in (t1, t2, t3, t4)):
            time.sleep(0.1)

def run_server(bind: str, pw: str, nick_bytes: bytes, stealth: bool):
    h, p = bind.split(":"); p = int(p)
    if not is_bind_ok(h):
        raise ValueError("bad listen address")
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as srv:
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind((h, p))
        srv.listen(1)
        sys.stdout.write(f"listening on {h}:{p}\n")
        sys.stdout.flush()
        conn, addr = srv.accept()
        with conn:
            set_keepalive(conn)
            peer = f"{addr[0]}:{addr[1]}"
            do_handshake_server(conn, pw, stealth=False)
            Chat(conn, peer, pw, nick_bytes, stealth=stealth).run()

def run_client(dst: str, pw: str, nick_bytes: bytes, stealth: bool):
    h, p = dst.split(":"); p = int(p)
    require_connect_ip(h)
    with socket.create_connection((h, p), timeout=10.0) as s:
        set_keepalive(s)
        peer = f"{h}:{p}"
        do_handshake_client(s, pw, stealth=False)
        Chat(s, peer, pw, nick_bytes, stealth=stealth).run()

def main():
    ap = argparse.ArgumentParser()
    m = ap.add_mutually_exclusive_group(required=True)
    m.add_argument("--listen",  help="IP:port to listen (use 0.0.0.0 or ::)")
    m.add_argument("--connect", help="IP:port to connect (global IP only)")
    ap.add_argument("--password", required=True)
    ap.add_argument("--nick",     help="ASCII nickname (printed as bytes)", default=None)
    ap.add_argument("--nick-hex", help="hex-encoded nickname bytes (overrides --nick)", default=None)
    ap.add_argument("--stealth",  action="store_true", help="enable stealth mode: padding, noise, hidden types, double PQS")
    a = ap.parse_args()

    sys.stdout.write((BANNER + b" | PQS " + getattr(PQS, "VERSION", b"?")).decode("ascii","ignore") + "\n")
    sys.stdout.flush()

    if a.nick_hex:
        try:
            nick_bytes = bytes.fromhex(a.nick_hex)
        except Exception:
            sys.stderr.write("bad --nick-hex\n"); sys.stderr.flush(); sys.exit(1)
    elif a.nick is not None:
        try:
            nick_bytes = a.nick.encode("ascii")
        except Exception:
            sys.stderr.write("nick must be ASCII or use --nick-hex\n"); sys.stderr.flush(); sys.exit(1)
    else:
        nick_bytes = b"peer"

    if a.listen:
        run_server(a.listen, a.password, nick_bytes, stealth=a.stealth)
    else:
        run_client(a.connect, a.password, nick_bytes, stealth=a.stealth)

if __name__ == "__main__":
    main()