#!/usr/bin/env python3
"""
NBD server v6 — chunk translation + backup GDT overlay.

The primary GDT has invalid checksums (written through PERC which stored
its own checksums). The backup GDT at block group 1 has valid checksums.
We serve the backup GDT bytes at the primary GDT location so the kernel
can validate and mount the filesystem.

Reads the backup GDT once at startup and caches it.
All other reads: pure chunk translation, no modification.

Usage:
    python3 nbd_server_v6.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
import sys

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

BSIZE         = 4096
GDT_ENTRY     = 64
BPG           = 32768
NUM_GROUPS    = 35728

# Primary GDT: virtual bytes 4096 to 4096+NUM_GROUPS*64
PRIMARY_GDT_START = BSIZE
PRIMARY_GDT_SIZE  = NUM_GROUPS * GDT_ENTRY
PRIMARY_GDT_END   = PRIMARY_GDT_START + PRIMARY_GDT_SIZE

# Backup GDT: at block group 1, block 1 = (BPG+1)*BSIZE
BACKUP_GDT_START  = (BPG + 1) * BSIZE

# Cached backup GDT (loaded at startup)
_backup_gdt = None

# NBD protocol
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 raw_read(virt_offset, length):
    """Pure chunk translation — no modifications."""
    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
            pos       += seg_len
            remaining -= seg_len
    return bytes(result)


def load_backup_gdt():
    """Read and cache the backup GDT from group 1."""
    global _backup_gdt
    print(f'[gdt] loading backup GDT from virtual byte {BACKUP_GDT_START}...')
    _backup_gdt = raw_read(BACKUP_GDT_START, PRIMARY_GDT_SIZE)

    # Verify first few entries look sane
    ok = True
    for i in range(min(5, NUM_GROUPS)):
        e = _backup_gdt[i*GDT_ENTRY:(i+1)*GDT_ENTRY]
        bb = struct.unpack_from('<I', e, 0)[0]
        ib = struct.unpack_from('<I', e, 4)[0]
        it = struct.unpack_from('<I', e, 8)[0]
        cs = struct.unpack_from('<H', e, 30)[0]
        print(f'[gdt]   group {i}: bb={bb} ib={ib} it={it} csum=0x{cs:04x}')
        if bb == 0 and ib == 0 and it == 0:
            ok = False

    if not ok:
        print('[gdt] WARNING: backup GDT looks empty, check parameters')
    else:
        print(f'[gdt] backup GDT loaded OK ({PRIMARY_GDT_SIZE//1024}KB)')
    return ok


def read_virtual(virt_offset, length):
    """
    Read with backup GDT overlay.
    Primary GDT region (virtual bytes 4096..4096+NUM_GROUPS*64) is
    served from the cached backup GDT instead of the primary location.
    Everything else is pure chunk translation.
    """
    req_end = virt_offset + length

    # Fast path: no overlap with primary GDT
    if req_end <= PRIMARY_GDT_START or virt_offset >= PRIMARY_GDT_END:
        return raw_read(virt_offset, length)

    # Build result from possibly multiple segments
    data = bytearray(raw_read(virt_offset, length))

    # Overlay backup GDT where request overlaps primary GDT
    ol_start = max(virt_offset, PRIMARY_GDT_START)
    ol_end   = min(req_end, PRIMARY_GDT_END)
    if ol_start < ol_end and _backup_gdt is not None:
        src_off = ol_start - PRIMARY_GDT_START
        dst_off = ol_start - virt_offset
        n       = ol_end - ol_start
        data[dst_off:dst_off + n] = _backup_gdt[src_off:src_off + n]

    return bytes(data)


def recv_all(conn, n):
    buf = b''
    while len(buf) < n:
        d = conn.recv(n - len(buf))
        if not d:
            raise ConnectionError('disconnected')
        buf += d
    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:
        conn.sendall(struct.pack('>Q', NBDMAGIC))
        conn.sendall(struct.pack('>Q', IHAVEOPT))
        conn.sendall(struct.pack('>H', 0x0003))
        recv_all(conn, 4)

        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:
                send_reply(conn, opt, NBD_REP_SERVER,
                           struct.pack('>I', 0))
                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 ({VIRT_SIZE//1024//1024//1024}GB)')

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

            if magic != NBD_REQUEST_MAGIC:
                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:
                return

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

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


def main():
    print('PERC H710 recovery NBD server v6')
    print(f'  device     : {DEV}')
    print(f'  lv_start   : byte {LV_PHYS_START}')
    print(f'  virt_size  : {VIRT_SIZE//1024//1024//1024} GB')
    print(f'  primary GDT: virtual bytes {PRIMARY_GDT_START}-{PRIMARY_GDT_END}')
    print(f'  backup GDT : virtual byte {BACKUP_GDT_START}')
    print()

    if not load_backup_gdt():
        print('ERROR: backup GDT load failed, check BACKUP_GDT_START')
        sys.exit(1)

    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()
    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()