#!/usr/bin/env python3
# B13W test client: RFB 3.8 + VeNCrypt X509Plain/TLSPlain (wayvnc enable_auth), raw 32bpp screenshot -> PPM, optional key/pointer input.
# vnc-vencrypt.py HOST PORT USER PASSFILE OUT.ppm [TEXT_TO_TYPE] [CERT_SHA256]
import socket, ssl, struct, sys, time, hashlib
h, p, user, pwf, out = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4], sys.argv[5]
text = sys.argv[6] if len(sys.argv) > 6 else ''; fp = sys.argv[7].replace(':', '').lower() if len(sys.argv) > 7 else ''
pw = open(pwf).read().strip()
s = socket.create_connection((h, p), 10)
def rx(n, c=None):
    c = c or s; b = b''
    while len(b) < n:
        d = c.recv(n - len(b))
        if not d: raise SystemExit('EOF')
        b += d
    return b
ver = rx(12); s.sendall(b'RFB 003.008\n')
n = rx(1)[0]
if n == 0: raise SystemExit('server refused: ' + rx(struct.unpack('>I', rx(4))[0]).decode())
types = list(rx(n)); print('security types:', types, 'None offered:', 1 in types)
assert 19 in types, 'no VeNCrypt'
s.sendall(b'\x13'); sv = rx(2); s.sendall(b'\x00\x02'); assert rx(1) == b'\x00'
m = rx(1)[0]; subs = [struct.unpack('>I', rx(4))[0] for _ in range(m)]; print('vencrypt subtypes:', subs)
st = 262 if 262 in subs else 259; s.sendall(struct.pack('>I', st)); assert rx(1) == b'\x01'
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT); ctx.check_hostname = False; ctx.verify_mode = ssl.CERT_NONE
t = ctx.wrap_socket(s); cert = t.getpeercert(binary_form=True)
cfp = hashlib.sha256(cert).hexdigest() if cert else '-'; print('tls', t.version(), t.cipher()[0], 'cert sha256', cfp[:16], 'pinned-match' if fp and cfp == fp else '')
if fp: assert cfp == fp, 'certificate fingerprint mismatch'
s = t
s.sendall(struct.pack('>II', len(user), len(pw)) + user.encode() + pw.encode())
r = struct.unpack('>I', rx(4))[0]
if r != 0:
    try: reason = rx(struct.unpack('>I', rx(4))[0]).decode()
    except SystemExit: reason = ''
    print('AUTH FAIL', reason); sys.exit(3)
print('AUTH OK subtype', st)
s.sendall(b'\x01'); w, hh = struct.unpack('>HH', rx(4)); rx(16); name = rx(struct.unpack('>I', rx(4))[0]).decode(); print('desktop', name, w, 'x', hh)
s.sendall(struct.pack('>BxxxBBBBHHHBBBxxx', 0, 32, 24, 0, 1, 255, 255, 255, 16, 8, 0)); s.sendall(struct.pack('>BxHi', 2, 1, 0))
def grab(path):
    s.sendall(struct.pack('>BBHHHH', 3, 0, 0, 0, w, hh)); img = bytearray(w * hh * 3); got = 0
    while got < w * hh:
        mt = rx(1)[0]
        if mt == 2: continue
        if mt == 1: rx(5); continue
        if mt != 0: raise SystemExit('msg %d' % mt)
        rx(1); nr = struct.unpack('>H', rx(2))[0]
        for _ in range(nr):
            x, y, rw, rh, enc = struct.unpack('>HHHHi', rx(12))
            if enc != 0: continue
            px = rx(rw * rh * 4)
            for j in range(rh):
                for i in range(rw):
                    v = px[(j*rw+i)*4:(j*rw+i)*4+4]; o = ((y+j)*w + x+i)*3; img[o:o+3] = bytes((v[2], v[1], v[0]))
            got += rw * rh
    open(path, 'wb').write(b'P6 %d %d 255\n' % (w, hh) + img); print('saved', path)
grab(out)
if text:
    s.sendall(struct.pack('>BBHH', 5, 0, 64, 64)); time.sleep(0.2)          # pointer to the centre (virtual pointer)
    for ch in text:
        k = ord(ch); s.sendall(struct.pack('>BBxxI', 4, 1, k)); s.sendall(struct.pack('>BBxxI', 4, 0, k)); time.sleep(0.05)
    time.sleep(1.5); grab(out.replace('.ppm', '-typed.ppm'))