#!/usr/bin/env python3
"""
NBD server v4 — fixed newstyle protocol size negotiation.

Key fix: NBD_OPT_GO info reply must send NBD_INFO_EXPORT (type 0) record
with correct format, followed by NBD_REP_ACK. Without this the client
connects but reports size=0.
"""

import socket
import struct
import threading

# ── Physical layout ───────────────────────────────────────────────────────────
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 block count

# ── ext4 filesystem parameters ────────────────────────────────────────────────
BSIZE      = 4096
GDT_ENTRY  = 64
NUM_GROUPS = 35728

GDT_START_VIRT = BSIZE
GDT_END_VIRT   = BSIZE + NUM_GROUPS * GDT_ENTRY

SB_VIRT_OFFSET = 1024
SB_SIZE        = 1024
SB_INCOMPAT_OFF  = 96
SB_RO_COMPAT_OFF = 100
SB_CHECKSUM_OFF  = 1020

INCOMPAT_HAS_JOURNAL    = 0x00000004
RO_COMPAT_METADATA_CSUM = 0x00000400
RO_COMPAT_GDT_CSUM      = 0x00000010

_patched_sb = None
_sb_lock    = threading.Lock()

# ── 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   # info type: export size + flags

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_WRITE = 1
NBD_CMD_DISC  = 2
NBD_CMD_FLUSH = 3

TRANSMISSION_FLAGS = NBD_FLAG_HAS_FLAGS | NBD_FLAG_READ_ONLY | NBD_FLAG_SEND_FLUSH


# ── Chunk translation ─────────────────────────────────────────────────────────

def raw_read(virt_offset, length):
    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 bytes(result)


# ── GDT synthesis ─────────────────────────────────────────────────────────────

def make_gdt_entry(n):
    gd = bytearray(64)
    struct.pack_into('<I', gd,  0, 1038 + n)       # block_bitmap_lo
    struct.pack_into('<I', gd,  4, 1054 + n)        # inode_bitmap_lo
    struct.pack_into('<I', gd,  8, 1070 + n * 512)  # inode_table_lo
    struct.pack_into('<I', gd, 32, 0)               # block_bitmap_hi = 0
    struct.pack_into('<I', gd, 36, 0)               # inode_bitmap_hi = 0
    struct.pack_into('<I', gd, 40, 0)               # inode_table_hi  = 0
    # bg_flags: set INODE_UNINIT and BLOCK_UNINIT for damaged groups
    # This tells ext4 to not validate bitmaps for these groups
    struct.pack_into('<H', gd, 18, 0x0003)          # EXT4_BG_INODE_UNINIT | EXT4_BG_BLOCK_UNINIT
    return bytes(gd)


def patch_gdt(data, virt_offset, length):
    pos = virt_offset
    remaining = length
    while remaining > 0:
        in_group  = pos % (5 * CHUNK_BYTES)
        chunk_idx = in_group // CHUNK_BYTES
        intra     = in_group % CHUNK_BYTES
        seg_len   = min(CHUNK_BYTES - intra, remaining)
        seg_end   = pos + seg_len
        if chunk_idx == 4:
            ol_start = max(pos, GDT_START_VIRT)
            ol_end   = min(seg_end, GDT_END_VIRT)
            if ol_start < ol_end:
                for byte_abs in range(ol_start, ol_end):
                    gdt_rel = byte_abs - GDT_START_VIRT
                    grp     = gdt_rel // GDT_ENTRY
                    byte_in = gdt_rel % GDT_ENTRY
                    if grp < NUM_GROUPS:
                        entry = make_gdt_entry(grp)
                        data[byte_abs - virt_offset] = entry[byte_in]
        pos       += seg_len
        remaining -= seg_len


# ── Superblock patch ──────────────────────────────────────────────────────────

def get_patched_sb():
    global _patched_sb
    with _sb_lock:
        if _patched_sb is not None:
            return _patched_sb
        sb = bytearray(raw_read(SB_VIRT_OFFSET, SB_SIZE))
        incompat  = struct.unpack_from('<I', sb, SB_INCOMPAT_OFF)[0]
        ro_compat = struct.unpack_from('<I', sb, SB_RO_COMPAT_OFF)[0]
        incompat  &= ~INCOMPAT_HAS_JOURNAL
        ro_compat &= ~(RO_COMPAT_METADATA_CSUM | RO_COMPAT_GDT_CSUM)
        struct.pack_into('<I', sb, SB_INCOMPAT_OFF,  incompat)
        struct.pack_into('<I', sb, SB_RO_COMPAT_OFF, ro_compat)
        struct.pack_into('<I', sb, SB_CHECKSUM_OFF,  0)
        # Add to get_patched_sb() in nbd_server_v4.py after the existing patches:

        # Clear checksum type (offset 222, 1 byte)
        sb[222] = 0

        # Clear checksum seed (offset 408, 4 bytes)
        struct.pack_into('<I', sb, 408, 0)

        _patched_sb = bytes(sb)
        print(f'[sb] patched incompat=0x{incompat:08x} ro_compat=0x{ro_compat:08x}')
        return _patched_sb


# ── Combined read ─────────────────────────────────────────────────────────────

def read_virtual(virt_offset, length):
    data = bytearray(raw_read(virt_offset, length))
    req_end = virt_offset + length

    # Patch superblock
    sb_e = SB_VIRT_OFFSET + SB_SIZE
    if virt_offset < sb_e and req_end > SB_VIRT_OFFSET:
        patched = get_patched_sb()
        cs = max(virt_offset, SB_VIRT_OFFSET) - virt_offset
        ce = min(req_end, sb_e) - virt_offset
        ss = max(virt_offset, SB_VIRT_OFFSET) - SB_VIRT_OFFSET
        data[cs:ce] = patched[ss:ss + (ce - cs)]

    # Patch GDT
    if virt_offset < GDT_END_VIRT and req_end > GDT_START_VIRT:
        patch_gdt(data, virt_offset, length)

    return bytes(data)


# ── NBD helpers ───────────────────────────────────────────────────────────────

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


def send_option_reply(conn, opt, reply_type, data=b''):
    """Send a structured option reply."""
    conn.sendall(struct.pack('>Q', REPLYMAGIC))
    conn.sendall(struct.pack('>I', opt))
    conn.sendall(struct.pack('>I', reply_type))
    conn.sendall(struct.pack('>I', len(data)))
    if data:
        conn.sendall(data)


def send_info_export(conn, opt):
    """
    Send NBD_REP_INFO with NBD_INFO_EXPORT record, then NBD_REP_ACK.
    This is what makes the client know the export size.

    NBD_INFO_EXPORT record layout:
      uint16  info_type = 0 (NBD_INFO_EXPORT)
      uint64  export_size
      uint16  transmission_flags
    """
    info_data = struct.pack('>HQH',
                            NBD_INFO_EXPORT,
                            VIRT_SIZE,
                            TRANSMISSION_FLAGS)
    send_option_reply(conn, opt, NBD_REP_INFO, info_data)
    send_option_reply(conn, opt, NBD_REP_ACK)


# ── Client handler ────────────────────────────────────────────────────────────

def handle_client(conn, addr):
    print(f'[nbd] connect from {addr}')
    try:
        # Server handshake: NBDMAGIC + IHAVEOPT + server flags
        conn.sendall(struct.pack('>Q', NBDMAGIC))
        conn.sendall(struct.pack('>Q', IHAVEOPT))
        # Server flags: FIXED_NEWSTYLE (bit 0) + NO_ZEROES (bit 1)
        conn.sendall(struct.pack('>H', 0x0003))

        # Client flags (4 bytes)
        recv_all(conn, 4)

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

            print(f'[nbd] {addr} opt={opt} len={opt_len}')

            if opt == NBD_OPT_EXPORT_NAME:
                # Old-style: send size + flags + (maybe) padding, no reply magic
                conn.sendall(struct.pack('>Q', VIRT_SIZE))
                conn.sendall(struct.pack('>H', TRANSMISSION_FLAGS))
                # NO_ZEROES flag set so skip 124-byte padding
                break

            elif opt == NBD_OPT_GO:
                # New-style: send NBD_REP_INFO then NBD_REP_ACK
                send_info_export(conn, opt)
                break

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

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

            else:
                send_option_reply(conn, opt, NBD_REP_ERR_UNSUP)

        print(f'[nbd] {addr} entering transmission, size={VIRT_SIZE}')

        # Transmission phase
        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={offset} len={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('PERC H710 recovery NBD server v4')
    print(f'  device    : {DEV}')
    print(f'  lv start  : byte {LV_PHYS_START}')
    print(f'  virt size : {VIRT_SIZE} bytes ({VIRT_SIZE//1024//1024//1024} GB)')
    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 -t ext4 -o ro,norecovery /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()