#!/usr/bin/env python3
"""
NBD server v8 - PERC H710 chunk translation + on-the-fly patches:
  1. Superblock: clear metadata_csum, gdt_csum, has_journal feature bits
  2. GDT: zero all checksum fields in every group descriptor

Nothing is written to disk. All patches applied in memory on reads.
"""
import socket, struct, threading

DEV           = '/dev/md0'
CHUNK_BYTES   = 128 * 512
LV_PHYS_START = 5120000 * 512
VIRT_SIZE     = 9365766144 * 512

# Superblock location and fields
SB_OFFSET     = 1024
SB_SIZE       = 1024
SB_COMPAT     = 92
SB_INCOMPAT   = 96
SB_RO_COMPAT  = 100
SB_JNLINUM    = 180
SB_CHECKSUM   = 1020
INCOMPAT_HAS_JOURNAL   = 0x004
COMPAT_HAS_JOURNAL     = 0x004
RO_COMPAT_GDT_CSUM     = 0x010
RO_COMPAT_METADATA_CSUM = 0x400

# GDT location and fields
GDT_OFFSET    = 4096          # block 1
GDT_ENTRY_SZ  = 64
NUM_GROUPS    = 35728
GDT_SIZE      = NUM_GROUPS * GDT_ENTRY_SZ
GDT_END       = GDT_OFFSET + GDT_SIZE
GDT_CSUM_OFF  = 30            # checksum offset within each entry

# 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 patching."""
    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 patch_superblock(data, req_start):
    """
    Patch superblock feature bits in a data buffer.
    req_start: virtual offset of data[0].
    """
    sb_start = SB_OFFSET
    sb_end   = SB_OFFSET + SB_SIZE
    req_end  = req_start + len(data)

    if req_start >= sb_end or req_end <= sb_start:
        return  # no overlap

    # Offsets into data buffer
    def patch_u32(sb_field_off, mask_clear):
        buf_off = sb_start + sb_field_off - req_start
        if 0 <= buf_off <= len(data) - 4:
            val = struct.unpack_from('<I', data, buf_off)[0]
            val &= ~mask_clear
            struct.pack_into('<I', data, buf_off, val)

    def zero_u32(sb_field_off):
        buf_off = sb_start + sb_field_off - req_start
        if 0 <= buf_off <= len(data) - 4:
            struct.pack_into('<I', data, buf_off, 0)

    patch_u32(SB_COMPAT,    COMPAT_HAS_JOURNAL)
    patch_u32(SB_INCOMPAT,  INCOMPAT_HAS_JOURNAL)
    patch_u32(SB_RO_COMPAT, RO_COMPAT_GDT_CSUM | RO_COMPAT_METADATA_CSUM)
    zero_u32(SB_JNLINUM)
    zero_u32(SB_CHECKSUM)


def patch_gdt(data, req_start):
    """
    Zero checksum field (offset 30) in every GDT entry that
    overlaps with this read buffer.
    req_start: virtual offset of data[0].
    """
    req_end = req_start + len(data)

    if req_start >= GDT_END or req_end <= GDT_OFFSET:
        return  # no overlap

    # First and last GDT entry indices that could overlap
    first_entry = max(0, (req_start - GDT_OFFSET) // GDT_ENTRY_SZ)
    last_entry  = min(NUM_GROUPS - 1,
                      (req_end - GDT_OFFSET - 1) // GDT_ENTRY_SZ)

    for g in range(first_entry, last_entry + 1):
        # Virtual offset of checksum field for group g
        csum_virt = GDT_OFFSET + g * GDT_ENTRY_SZ + GDT_CSUM_OFF
        buf_off   = csum_virt - req_start
        if 0 <= buf_off <= len(data) - 2:
            struct.pack_into('<H', data, buf_off, 0)


def read_virtual(virt_offset, length):
    data = raw_read(virt_offset, length)
    patch_superblock(data, virt_offset)
    patch_gdt(data, virt_offset)
    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} in 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 v8')
    print(f'  device    : {DEV}')
    print(f'  lv_start  : byte {LV_PHYS_START}')
    print(f'  virt_size : {VIRT_SIZE // 1024**3} GB')
    print(f'  patches   : superblock features + GDT checksums (on-the-fly)')
    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('  fls /dev/nbd0 1585918')
    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()
