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

No feature flag patching. No checksum zeroing. No journal disabling.
Just serves the filesystem exactly as it was written, with a complete
GDT built by merging primary (for non-metadata-chunk entries) and
backup group 1 (for entries that fall in metadata chunks).

Primary and backup GDT bad entries are completely disjoint (0 overlap),
so this gives 100% GDT coverage with authentic data and valid checksums.

Usage:
    # First build the merged GDT:
    # python3 build_merged_gdt.py  (saves /tmp/merged_gdt.bin)
    
    python3 nbd_server_v9.py &
    nbd-client 127.0.0.1 10809 /dev/nbd0 -N ""
    ext4magic /dev/nbd0 -s 4096 -n 32768 -R -I 1585918 -a $(date -d "2023-01-01" +%s) -d /mnt/recovered
"""

import socket, struct, threading, sys, os

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

# GDT location in virtual address space
BSIZE         = 4096
GDT_ENTRY_SZ  = 64
NUM_GROUPS    = 35728
GDT_VIRT_START = BSIZE                          # block 1 = byte 4096
GDT_VIRT_END   = BSIZE + NUM_GROUPS * GDT_ENTRY_SZ

# 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

# 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.')
    print('Run the GDT builder script 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, '
          f'expected {expected}')
    sys.exit(1)

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


def raw_read(virt_offset, length):
    """Pure PERC 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)
                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.
    Only the primary GDT region (bytes 4096 to 4096+NUM_GROUPS*64)
    is modified — replaced with the pre-built merged GDT.
    Everything else is pure chunk translation, unmodified.
    """
    data    = raw_read(virt_offset, length)
    req_end = virt_offset + length

    # Overlay merged GDT where request overlaps primary GDT region
    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 v9')
    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'  patch      : merged GDT only (primary + backup group 1)')
    print(f'  no patches : superblock, features, checksums all authentic')
    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()
    print('Then try:')
    print('  e2fsck -n /dev/nbd0')
    print('  fls /dev/nbd0 1585918')
    print('  ext4magic /dev/nbd0 -s 4096 -n 32768 -R -I 1585918 \\')
    print('    -a $(date -d "2023-01-01" +%s) -d /mnt/recovered')
    print()

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


if __name__ == '__main__':
    main()
