#!/usr/bin/env python3
"""
NBD server v5 — minimal PERC H710 chunk translation only.

Simply reproduces what the PERC controller presented to the OS:
  - Skip every 5th 64KB chunk (PERC internal metadata)
  - Serve the result as a read-only block device

No superblock patching. No GDT synthesis. The filesystem was written
correctly through the PERC's translation — reading it back the same
way should give a valid filesystem.

Usage:
    python3 nbd_server_v5.py &
    nbd-client 127.0.0.1 10809 /dev/nbd0 -N ""
    mount -o ro,norecovery -t ext4 /dev/nbd0 /mnt/root
"""

import socket
import struct
import threading

DEV           = '/dev/md0'
CHUNK_BYTES   = 128 * 512          # 64 KB per chunk
LV_PHYS_START = 5120000 * 512      # byte 2,621,440,000
VIRT_SIZE     = 9365766144 * 512   # from superblock block count

# NBD newstyle protocol constants
NBDMAGIC          = 0x4e42444d41474943
IHAVEOPT          = 0x49484156454F5054
REPLYMAGIC        = 0x3e889045565a9
NBD_OPT_EXPORT_NAME = 1
NBD_OPT_ABORT       = 2
NBD_OPT_LIST        = 3
NBD_OPT_GO          = 7
NBD_REP_ACK         = 1
NBD_REP_SERVER      = 2
NBD_REP_INFO        = 3
NBD_REP_ERR_UNSUP   = (1 << 31) | 1
NBD_INFO_EXPORT     = 0
NBD_FLAG_HAS_FLAGS  = 1 << 0
NBD_FLAG_READ_ONLY  = 1 << 1
NBD_FLAG_SEND_FLUSH = 1 << 2
NBD_REQUEST_MAGIC   = 0x25609513
NBD_REPLY_MAGIC     = 0x67446698
NBD_CMD_READ        = 0
NBD_CMD_DISC        = 2
NBD_CMD_FLUSH       = 3

TX_FLAGS = NBD_FLAG_HAS_FLAGS | NBD_FLAG_READ_ONLY | NBD_FLAG_SEND_FLUSH


def read_virtual(virt_offset, length):
    """
    Read length bytes from the virtual address space.
    Applies PERC chunk translation: every 5th 64KB chunk is skipped
    (it contained PERC internal metadata and is not part of user data).
    Skipped chunks return zeros.
    """
    result = bytearray(length)
    pos = virt_offset
    remaining = length

    with open(DEV, 'rb') as f:
        while remaining > 0:
            group     = pos // (5 * CHUNK_BYTES)
            in_group  = pos % (5 * CHUNK_BYTES)
            chunk_idx = in_group // CHUNK_BYTES
            intra     = in_group % CHUNK_BYTES
            seg_len   = min(CHUNK_BYTES - intra, remaining)
            dst_off   = pos - virt_offset

            if chunk_idx != 4:
                phys = (LV_PHYS_START
                        + group * 4 * CHUNK_BYTES
                        + chunk_idx * CHUNK_BYTES
                        + intra)
                f.seek(phys)
                chunk = f.read(seg_len)
                result[dst_off:dst_off + len(chunk)] = chunk
            # chunk_idx == 4: leave as zeros (PERC metadata, not user data)

            pos       += seg_len
            remaining -= seg_len

    return bytes(result)


def recv_all(conn, n):
    buf = b''
    while len(buf) < n:
        data = conn.recv(n - len(buf))
        if not data:
            raise ConnectionError('client disconnected')
        buf += data
    return buf


def send_reply(conn, opt, rtype, data=b''):
    conn.sendall(struct.pack('>Q', REPLYMAGIC))
    conn.sendall(struct.pack('>I', opt))
    conn.sendall(struct.pack('>I', rtype))
    conn.sendall(struct.pack('>I', len(data)))
    if data:
        conn.sendall(data)


def handle_client(conn, addr):
    print(f'[nbd] {addr} connected')
    try:
        # Handshake
        conn.sendall(struct.pack('>Q', NBDMAGIC))
        conn.sendall(struct.pack('>Q', IHAVEOPT))
        conn.sendall(struct.pack('>H', 0x0003))  # FIXED_NEWSTYLE | NO_ZEROES
        recv_all(conn, 4)  # client flags

        # Option haggling
        while True:
            hdr = recv_all(conn, 16)
            _, opt, opt_len = struct.unpack('>QII', hdr)
            opt_data = recv_all(conn, opt_len) if opt_len else b''

            if opt == NBD_OPT_EXPORT_NAME:
                conn.sendall(struct.pack('>Q', VIRT_SIZE))
                conn.sendall(struct.pack('>H', TX_FLAGS))
                break

            elif opt == NBD_OPT_GO:
                info = struct.pack('>HQH', NBD_INFO_EXPORT, VIRT_SIZE, TX_FLAGS)
                send_reply(conn, opt, NBD_REP_INFO, info)
                send_reply(conn, opt, NBD_REP_ACK)
                break

            elif opt == NBD_OPT_LIST:
                name = b''
                send_reply(conn, opt, NBD_REP_SERVER,
                           struct.pack('>I', 0) + name)
                send_reply(conn, opt, NBD_REP_ACK)

            elif opt == NBD_OPT_ABORT:
                send_reply(conn, opt, NBD_REP_ACK)
                return

            else:
                send_reply(conn, opt, NBD_REP_ERR_UNSUP)

        print(f'[nbd] {addr} transmission phase ({VIRT_SIZE//1024//1024//1024}GB)')

        # Transmission
        while True:
            hdr = recv_all(conn, 28)
            magic, flags, cmd, handle, offset, length = \
                struct.unpack('>IHHQQI', hdr)

            if magic != NBD_REQUEST_MAGIC:
                print(f'[nbd] bad magic 0x{magic:08x}')
                return

            if cmd == NBD_CMD_READ:
                try:
                    payload = read_virtual(offset, length)
                except Exception as e:
                    print(f'[nbd] read error @ {offset}+{length}: {e}')
                    payload = b'\x00' * length
                conn.sendall(struct.pack('>IIQ', NBD_REPLY_MAGIC, 0, handle))
                conn.sendall(payload)

            elif cmd == NBD_CMD_FLUSH:
                conn.sendall(struct.pack('>IIQ', NBD_REPLY_MAGIC, 0, handle))

            elif cmd == NBD_CMD_DISC:
                print(f'[nbd] {addr} disconnect')
                return

            else:
                conn.sendall(struct.pack('>IIQ', NBD_REPLY_MAGIC, 1, handle))

    except (ConnectionError, BrokenPipeError, ConnectionResetError):
        print(f'[nbd] {addr} dropped')
    except Exception as e:
        print(f'[nbd] {addr} error: {e}')
        import traceback; traceback.print_exc()
    finally:
        conn.close()


def main():
    print(f'PERC H710 recovery NBD server v5 (minimal)')
    print(f'  device    : {DEV}')
    print(f'  lv_start  : byte {LV_PHYS_START} (sector {LV_PHYS_START//512})')
    print(f'  virt_size : {VIRT_SIZE//1024//1024//1024} GB')
    print(f'  chunk     : {CHUNK_BYTES//1024} KB, every 5th skipped')
    print()

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(('127.0.0.1', 10809))
    srv.listen(5)

    print('Listening on 127.0.0.1:10809')
    print('  nbd-client 127.0.0.1 10809 /dev/nbd0 -N ""')
    print('  mount -o ro,norecovery -t ext4 /dev/nbd0 /mnt/root')
    print()

    while True:
        conn, addr = srv.accept()
        threading.Thread(target=handle_client, args=(conn, addr),
                         daemon=True).start()


if __name__ == '__main__':
    main()
