#!/usr/bin/env python3
"""
NBD server v10 - PERC H710 corrected chunk translation + merged GDT overlay.

Key fix over v9: all 5 chunks per stripe group contain filesystem data.
The old formula used a 4-chunk stride between groups, making ~20% of blocks
unreadable and shifting every address past group 0 by an accumulating error.

Correct formula: phys = LV_START + group*5*CHUNK + chunk_idx*CHUNK + intra
for ALL chunk_idx values 0-4.

Usage:
    python3 build_merged_v2.py        # rebuild /tmp/merged_gdt.bin
    python3 nbd_server_v10.py &
    nbd-client 127.0.0.1 10809 /dev/nbd0 -N ""
    mount -o ro /dev/nbd0 /mnt/fs
"""

import socket, struct, threading, sys, os

DEV           = '/dev/md0'
MERGED_GDT    = '/tmp/merged_gdt.bin'
CHUNK_BYTES   = 128 * 512          # 64 KB
LV_PHYS_START = 5120000 * 512
VIRT_SIZE     = 9365766144 * 512

BSIZE         = 4096
GDT_ENTRY_SZ  = 64
NUM_GROUPS    = 35728
GDT_VIRT_START = BSIZE
GDT_VIRT_END   = BSIZE + NUM_GROUPS * GDT_ENTRY_SZ

# NBD 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

# Load merged GDT at startup
print(f'Loading merged GDT from {MERGED_GDT}...')
if not os.path.exists(MERGED_GDT):
    print(f'ERROR: {MERGED_GDT} not found. Run build_merged_v2.py first.')
    sys.exit(1)

with open(MERGED_GDT, 'rb') as f:
    MERGED_GDT_DATA = f.read()

expected = NUM_GROUPS * GDT_ENTRY_SZ
if len(MERGED_GDT_DATA) != expected:
    print(f'ERROR: merged GDT is {len(MERGED_GDT_DATA)} bytes, expected {expected}')
    sys.exit(1)

print(f'Merged GDT loaded: {len(MERGED_GDT_DATA)//1024}KB ({NUM_GROUPS} groups)')


def raw_read(virt_offset, length):
    """PERC H710 chunk translation — all 5 chunks per group are filesystem data."""
    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
            phys = (LV_PHYS_START
                    + group * 5 * CHUNK_BYTES
                    + chunk_idx * CHUNK_BYTES
                    + intra)
            f.seek(phys)
            data = f.read(seg_len)
            result[dst_off:dst_off + len(data)] = data
            pos       += seg_len
            remaining -= seg_len
    return result


def read_virtual(virt_offset, length):
    """Read with merged GDT overlay over the primary GDT region only."""
    data    = raw_read(virt_offset, length)
    req_end = virt_offset + length

    if virt_offset < GDT_VIRT_END and req_end > GDT_VIRT_START:
        ol_start = max(virt_offset, GDT_VIRT_START)
        ol_end   = min(req_end, GDT_VIRT_END)
        src_off  = ol_start - GDT_VIRT_START
        dst_off  = ol_start - virt_offset
        n        = ol_end - ol_start
        data[dst_off:dst_off + n] = MERGED_GDT_DATA[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')
        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:
                print(f'[nbd] {addr} disconnected')
                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('PERC H710 recovery NBD server v10')
    print(f'  device     : {DEV}')
    print(f'  lv_start   : byte {LV_PHYS_START}')
    print(f'  virt_size  : {VIRT_SIZE // 1024**3} GB')
    print(f'  GDT region : bytes {GDT_VIRT_START}-{GDT_VIRT_END}')
    print(f'  translation: group*5*CHUNK (all 5 chunks are filesystem data)')
    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()

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


if __name__ == '__main__':
    main()
