#!/usr/bin/env python3
"""A scriptable, hostile 9P2000 server on a Unix socket (stdlib only).

Usage: adv_bridge_hostile.py SOCKET MODE

Serves a tiny in-memory tree:
    /f          "hello world\\n"
    /d/g        "in d\\n"
    /fids       reading it returns the number of fids currently bound
    /big        1 MiB of pseudo-random bytes
plus create/write/remove/wstat so the scratch battery can run in `ok` mode.

MODE selects one misbehaviour (see MODES below). Everything not covered by
the mode behaves normally, so 9ns gets through version/attach/stat(root).
"""
import os
import random
import socket
import struct
import sys
import time

NOTAG = 0xFFFF
NOFID = 0xFFFFFFFF
QTDIR = 0x80
DMDIR = 0x80000000

Tversion, Rversion = 100, 101
Tauth, Rauth = 102, 103
Tattach, Rattach = 104, 105
Rerror = 107
Tflush, Rflush = 108, 109
Twalk, Rwalk = 110, 111
Topen, Ropen = 112, 113
Tcreate, Rcreate = 114, 115
Tread, Rread = 116, 117
Twrite, Rwrite = 118, 119
Tclunk, Rclunk = 120, 121
Tremove, Rremove = 122, 123
Tstat, Rstat = 124, 125
Twstat, Rwstat = 126, 127

MODES = """
ok              behave (qid paths are recycled LIFO after remove, like many servers)
trunc           Rread on /f: send half the frame, then close
short_frame     Rread on /f: frame whose size field is 3
huge_frame      Rread on /f: frame whose size field is msize+1
wrong_tag       Rread on /f: reply carries tag+1
wrong_type      Tstat on /f: answer with an Rwalk
rread_big       Rread on /f: count = requested+1
rwalk_many      Twalk to f: nwqid = nwname+1
rwalk_zero      Twalk to nope: Rwalk nwqid=0 instead of Rerror
rstat_garbage   Tstat on /f: random bytes as the stat
rstat_overlong  Tstat on /f: inner stat size disagrees with outer
dir_split       Tread on /: a stat record split across two Rreads
dir_forever     Tread on /: ignore offset, always return the same records
qid_collide     every file and dir shares qid.path 7 (root keeps its own)
qid_zero        every qid.path is 0, including the root
name_slash      / has an entry "a/b"
name_empty      / has an entry ""
name_huge       / has an entry with a 60000-byte name
name_dots       / lists "." and ".." too
rerror_big      Twalk to nope: Rerror with 65535 bytes of text
extra_reply     Rread on /f: an unsolicited Rclunk (tag 9) precedes the real reply
never           Tread on /f: never reply (hang); the server stops reading, so a Tflush is never seen
never_flush     Tread on /f: never reply, but keep serving: log every Tflush and answer Rflush
close_mid       Tread on /f: close the socket without replying
renegotiate     Tread on /f: an unsolicited Rversion precedes the real reply
length_max      Tstat on /f: length = 2**64-1
iounit_one      Ropen: iounit = 1
rwrite_big      Rwrite: count = requested+1
msize_tiny      Rversion msize = 64
version_unknown Rversion "unknown"
rename_fail     Twstat with a new name always fails "file already exists"
slow            every reply delayed 20 ms (for interrupt tests)
"""


def s8(x): return struct.pack('<B', x)
def s16(x): return struct.pack('<H', x)
def s32(x): return struct.pack('<I', x)
def s64(x): return struct.pack('<Q', x)
def sstr(b):
    if isinstance(b, str):
        b = b.encode()
    return s16(len(b)) + b


class Node:
    def __init__(self, name, isdir, path, content=b''):
        self.name = name
        self.isdir = isdir
        self.path = path
        self.content = bytearray(content)
        self.children = {}
        self.mode = 0o755 if isdir else 0o644
        self.mtime = int(time.time())
        self.removed = False

    def qid(self, srv):
        path = self.path
        if srv.mode == 'qid_zero':
            path = 0
        elif srv.mode == 'qid_collide' and self is not srv.root:
            path = 7
        return s8(QTDIR if self.isdir else 0) + s32(1) + s64(path)

    def stat_bytes(self, srv, name=None, length=None):
        if name is None:
            name = self.name
        if length is None:
            length = 0 if self.isdir else len(self.content)
        body = (s16(0) + s32(0) + self.qid(srv) + s32((DMDIR if self.isdir else 0) | self.mode)
                + s32(self.mtime) + s32(self.mtime) + s64(length)
                + sstr(name) + sstr('u') + sstr('g') + sstr('u'))
        return s16(len(body)) + body


class Server:
    def __init__(self, mode):
        self.mode = mode
        self.next_path = 100
        self.free_paths = []
        self.root = Node('', True, 1)
        f = self.mk(self.root, 'f', False, b'hello world\n')
        d = self.mk(self.root, 'd', True)
        self.mk(d, 'g', False, b'in d\n')
        self.mk(self.root, 'fids', False)
        rnd = random.Random(4)
        self.mk(self.root, 'big', False, bytes(rnd.getrandbits(8) for _ in range(1 << 20)))
        self.fids = {}
        self.msize = 8192
        self.ops = 0

    def alloc_path(self):
        if self.free_paths:
            return self.free_paths.pop()
        self.next_path += 1
        return self.next_path

    def mk(self, parent, name, isdir, content=b''):
        n = Node(name, isdir, self.alloc_path(), content)
        parent.children[name] = n
        return n

    def log(self, text):
        print(text, flush=True)

    # -- framing ---------------------------------------------------------
    def frame(self, typ, tag, body):
        return s32(7 + len(body)) + s8(typ) + s16(tag) + body

    def err(self, tag, text):
        return self.frame(Rerror, tag, sstr(text))

    def serve(self, conn):
        buf = b''
        while True:
            if len(buf) >= 4:
                n = struct.unpack('<I', buf[:4])[0]
                if len(buf) >= n:
                    msg, buf = buf[:n], buf[n:]
                    out = self.handle(msg)
                    if out is None:
                        return  # hang up / hang
                    if self.mode == 'slow':
                        time.sleep(0.02)
                    conn.sendall(out)
                    continue
            data = conn.recv(65536)
            if not data:
                return
            buf += data

    def handle(self, msg):
        typ = msg[4]
        tag = struct.unpack('<H', msg[5:7])[0]
        b = msg[7:]
        self.ops += 1
        r = Reader(b)
        if typ == Tversion:
            msize = r.u32()
            ver = r.str()
            self.msize = min(msize, 1 << 20)
            self.fids = {}
            if self.mode == 'msize_tiny':
                self.msize = 64
            if self.mode == 'version_unknown':
                return self.frame(Rversion, NOTAG, s32(self.msize) + sstr('unknown'))
            return self.frame(Rversion, NOTAG, s32(self.msize) + sstr('9P2000'))
        if typ == Tauth:
            return self.err(tag, 'authentication not required')
        if typ == Tattach:
            fid = r.u32()
            r.u32()
            r.str()
            r.str()
            if fid in self.fids:
                return self.err(tag, 'fid in use')
            self.fids[fid] = [self.root, False]
            return self.frame(Rattach, tag, self.root.qid(self))
        if typ == Tflush:
            oldtag = r.u16()
            self.log('Tflush tag=%d oldtag=%d' % (tag, oldtag))
            return self.frame(Rflush, tag, b'')
        if typ == Twalk:
            fid, newfid, nw = r.u32(), r.u32(), r.u16()
            names = [r.str() for _ in range(nw)]
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            if newfid != fid and newfid in self.fids:
                return self.err(tag, 'fid in use')
            node = self.fids[fid][0]
            qids = b''
            n = 0
            for name in names:
                if not node.isdir:
                    break
                if name == '..':
                    nxt = node  # root's parent is itself; good enough
                elif name in node.children:
                    nxt = node.children[name]
                else:
                    break
                node = nxt
                qids += node.qid(self)
                n += 1
            if n < nw and n == 0:
                if self.mode == 'rwalk_zero' and names == ['nope']:
                    return self.frame(Rwalk, tag, s16(0))
                if self.mode == 'rerror_big' and names == ['nope']:
                    return self.err(tag, 'x' * 65535)
                return self.err(tag, 'file does not exist')
            if n == nw:
                self.fids[newfid] = [node, False]
            if self.mode == 'rwalk_many' and names == ['f']:
                return self.frame(Rwalk, tag, s16(n + 1) + qids + node.qid(self))
            return self.frame(Rwalk, tag, s16(n) + qids)
        if typ == Topen:
            fid, mode = r.u32(), r.u8()
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            node = self.fids[fid][0]
            if node.isdir and (mode & 3) != 0:
                return self.err(tag, 'is a directory')
            if mode & 0x10 and not node.isdir:
                node.content = bytearray()
            self.fids[fid][1] = True
            iounit = 1 if self.mode == 'iounit_one' else 0
            return self.frame(Ropen, tag, node.qid(self) + s32(iounit))
        if typ == Tcreate:
            fid = r.u32()
            name = r.str()
            perm = r.u32()
            mode = r.u8()
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            parent = self.fids[fid][0]
            if not parent.isdir:
                return self.err(tag, 'not a directory')
            if name in parent.children:
                return self.err(tag, 'file already exists')
            node = self.mk(parent, name, bool(perm & DMDIR))
            node.mode = perm & 0o777
            self.fids[fid] = [node, True]
            return self.frame(Rcreate, tag, node.qid(self) + s32(0))
        if typ == Tread:
            fid, off, count = r.u32(), r.u64(), r.u32()
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            node, opened = self.fids[fid]
            if not opened:
                return self.err(tag, 'fid not open')
            if node.isdir:
                return self.readdir(tag, node, off, count)
            if node.name == 'fids':
                data = ('%d\n' % len(self.fids)).encode()
                data = data[off:off + count]
                return self.frame(Rread, tag, s32(len(data)) + data)
            data = bytes(node.content[off:off + count])
            if node is self.root.children.get('f'):
                m = self.mode
                if m == 'trunc':
                    fr = self.frame(Rread, tag, s32(len(data)) + data)
                    self.conn.sendall(fr[:len(fr) // 2])
                    return None
                if m == 'short_frame':
                    return s32(3) + s8(Rread) + s16(tag)
                if m == 'huge_frame':
                    return s32(self.msize + 1) + s8(Rread) + s16(tag) + s32(len(data)) + data
                if m == 'wrong_tag':
                    return self.frame(Rread, (tag + 1) & 0xFFFF, s32(len(data)) + data)
                if m == 'rread_big':
                    data = b'x' * (count + 1)
                    return self.frame(Rread, tag, s32(len(data)) + data)
                if m == 'extra_reply':
                    return self.frame(Rclunk, 9, b'') + self.frame(Rread, tag, s32(len(data)) + data)
                if m == 'never':
                    time.sleep(3600)
                    return None
                if m == 'never_flush':
                    self.log('Tread tag=%d fid=%d hang' % (tag, fid))
                    return b''
                if m == 'close_mid':
                    return None
                if m == 'renegotiate':
                    return self.frame(Rversion, NOTAG, s32(self.msize) + sstr('9P2000')) + self.frame(Rread, tag, s32(len(data)) + data)
            return self.frame(Rread, tag, s32(len(data)) + data)
        if typ == Twrite:
            fid, off = r.u32(), r.u64()
            data = r.data()
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            node, opened = self.fids[fid]
            if not opened or node.isdir:
                return self.err(tag, 'fid not open for write')
            if node.name == 'fids':
                return self.err(tag, 'permission denied')
            if off > len(node.content):
                node.content.extend(b'\0' * (off - len(node.content)))
            node.content[off:off + len(data)] = data
            node.mtime = int(time.time())
            n = len(data) + 1 if self.mode == 'rwrite_big' else len(data)
            return self.frame(Rwrite, tag, s32(n))
        if typ == Tclunk:
            fid = r.u32()
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            del self.fids[fid]
            return self.frame(Rclunk, tag, b'')
        if typ == Tremove:
            fid = r.u32()
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            node = self.fids[fid][0]
            del self.fids[fid]
            if node is self.root:
                return self.err(tag, 'cannot remove root')
            if node.isdir and node.children:
                return self.err(tag, 'directory not empty')
            parent = self.find_parent(self.root, node)
            if parent is not None:
                del parent.children[node.name]
                self.free_paths.append(node.path)
            node.removed = True
            return self.frame(Rremove, tag, b'')
        if typ == Tstat:
            fid = r.u32()
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            node = self.fids[fid][0]
            if node is self.root.children.get('f'):
                m = self.mode
                if m == 'wrong_type':
                    return self.frame(Rwalk, tag, s16(0))
                if m == 'rstat_garbage':
                    junk = bytes([0xAB] * 60)
                    return self.frame(Rstat, tag, s16(len(junk)) + junk)
                if m == 'rstat_overlong':
                    st = node.stat_bytes(self)
                    inner = st[2:]
                    return self.frame(Rstat, tag, s16(len(inner) + 5) + inner)
                if m == 'length_max':
                    st = node.stat_bytes(self, length=2 ** 64 - 1)
                    return self.frame(Rstat, tag, s16(len(st)) + st)
            st = node.stat_bytes(self)
            return self.frame(Rstat, tag, s16(len(st)) + st)
        if typ == Twstat:
            fid = r.u32()
            r.u16()
            st = r.bytes(r.u16())
            if fid not in self.fids:
                return self.err(tag, 'unknown fid')
            node = self.fids[fid][0]
            sr = Reader(st)
            sr.u16(); sr.u32(); sr.bytes(13)
            mode = sr.u32(); sr.u32(); mtime = sr.u32(); length = sr.u64()
            name = sr.str()
            if name and name != node.name:
                if self.mode == 'rename_fail':
                    return self.err(tag, 'file already exists')
                parent = self.find_parent(self.root, node)
                if name in parent.children:
                    return self.err(tag, 'file already exists')
                del parent.children[node.name]
                node.name = name
                parent.children[name] = node
            if mode != 0xFFFFFFFF:
                node.mode = mode & 0o777
            if mtime != 0xFFFFFFFF:
                node.mtime = mtime
            if length != 0xFFFFFFFFFFFFFFFF and not node.isdir:
                if length < len(node.content):
                    del node.content[length:]
                else:
                    node.content.extend(b'\0' * (length - len(node.content)))
            return self.frame(Rwstat, tag, b'')
        return self.err(tag, 'unsupported message')

    def find_parent(self, cur, node):
        for c in cur.children.values():
            if c is node:
                return cur
            if c.isdir:
                p = self.find_parent(c, node)
                if p is not None:
                    return p
        return None

    def readdir(self, tag, node, off, count):
        recs = []
        if node is self.root:
            m = self.mode
            if m == 'name_slash':
                recs.append(node.stat_bytes(self, name='a/b'))
            if m == 'name_empty':
                recs.append(node.stat_bytes(self, name=''))
            if m == 'name_huge':
                recs.append(node.stat_bytes(self, name='h' * 60000))
            if m == 'name_dots':
                recs.append(node.stat_bytes(self, name='.'))
                recs.append(node.stat_bytes(self, name='..'))
        for c in node.children.values():
            recs.append(c.stat_bytes(self))
        blob = b''.join(recs)
        if node is self.root and self.mode == 'dir_forever':
            return self.frame(Rread, tag, s32(len(blob)) + blob)
        if node is self.root and self.mode == 'dir_split':
            # first read: up to the middle of the second record; second read: the rest
            cut = len(recs[0]) + len(recs[1]) // 2
            if off == 0:
                data = blob[:cut]
            elif off == cut:
                data = blob[cut:]
            else:
                data = b''
            return self.frame(Rread, tag, s32(len(data)) + data)
        # 9P rule: offset 0 or previous offset+count; never split a record.
        out = b''
        pos = 0
        for rec in recs:
            if pos >= off and len(out) + len(rec) <= count:
                out += rec
            elif pos >= off:
                break
            pos += len(rec)
        return self.frame(Rread, tag, s32(len(out)) + out)


class Reader:
    def __init__(self, b):
        self.b = b
        self.i = 0

    def bytes(self, n):
        v = self.b[self.i:self.i + n]
        self.i += n
        return v

    def u8(self): return struct.unpack('<B', self.bytes(1))[0]
    def u16(self): return struct.unpack('<H', self.bytes(2))[0]
    def u32(self): return struct.unpack('<I', self.bytes(4))[0]
    def u64(self): return struct.unpack('<Q', self.bytes(8))[0]
    def str(self): return self.bytes(self.u16()).decode('utf-8', 'surrogateescape')
    def data(self): return self.bytes(self.u32())


def main():
    if len(sys.argv) != 3:
        print(__doc__ + MODES)
        sys.exit(2)
    path, mode = sys.argv[1], sys.argv[2]
    if mode not in [l.split()[0] for l in MODES.strip().splitlines()]:
        print('unknown mode', mode)
        sys.exit(2)
    try:
        os.unlink(path)
    except FileNotFoundError:
        pass
    ls = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    ls.bind(path)
    ls.listen(8)
    while True:
        conn, _ = ls.accept()
        srv = Server(mode)
        srv.conn = conn
        try:
            srv.serve(conn)
        except (BrokenPipeError, ConnectionResetError):
            pass
        finally:
            conn.close()


if __name__ == '__main__':
    main()