#!/usr/bin/env python3
"""
NBD server for PERC H710 RAID recovery.

Applies three transformations on every read:
  1. Chunk translation  — skips every 5th 64KB chunk (PERC internal metadata)
  2. Superblock patch   — clears metadata_csum and has_journal feature bits
                          so the kernel stops validating checksums we can't fix
  3. GDT reconstruction — synthesizes correct group descriptor entries for
                          regions that fall inside metadata chunks (zeros)

Confirmed filesystem parameters (from session forensics):
  bpg        = 32768   blocks per group
  ipg        = 8192    inodes per group
  inode_size = 256     bytes
  GDT entry  = 64      bytes (64bit feature)
  num_groups = 35728
  Group N:   bb = 1038+N,  ib = 1054+N,  it = 1070 + N*512

Usage:
    python3 nbd_server_v2.py &
    nbd-client -g 127.0.0.1 10809 /dev/nbd0
    mount -o ro,norecovery -t ext4 /dev/nbd0 /mnt/root
"""

import socket
import struct
import threading
import os
import sys

# ── 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               # block size in bytes
BPG           = 32768              # blocks per group
IPG           = 8192               # inodes per group
INODE_SIZE    = 256                # inode size in bytes
GDT_ENTRY     = 64                 # group descriptor size (64-bit mode)
NUM_GROUPS    = 35728              # total number of block groups

# GDT starts at block 1 = byte 4096 from LV start
GDT_START_VIRT = BSIZE             # virtual byte offset of GDT
GDT_END_VIRT   = BSIZE + NUM_GROUPS * GDT_ENTRY  # ~9.1 MB

# Superblock is at virtual byte 1024
SB_VIRT_OFFSET = 1024
SB_SIZE        = 1024

# Feature flag offsets within superblock
SB_FEAT_COMPAT_OFF    = 92   # s_feature_compat
SB_FEAT_INCOMPAT_OFF  = 96   # s_feature_incompat
SB_FEAT_RO_COMPAT_OFF = 100  # s_feature_ro_compat
SB_CHECKSUM_OFF       = 1020 # s_checksum (last 4 bytes of sb)

# Bits to clear
INCOMPAT_HAS_JOURNAL  = 0x00000004
RO_COMPAT_METADATA_CSUM = 0x00000400
RO_COMPAT_GDT_CSUM    = 0x00000010

# Cached patched superblock (built once on first read)
_patched_sb = None
_sb_lock = threading.Lock()

# ── NBD protocol constants ────────────────────────────────────────────────────
NBD_MAGIC         = 0x4e42444d41474943
NBD_CLISERV_MAGIC = 0x00420281861253
NBD_REQUEST_MAGIC = 0x25609513
NBD_REPLY_MAGIC   = 0x67446698
NBD_CMD_READ  = 0
NBD_CMD_WRITE = 1
NBD_CMD_DISC  = 2
NBD_FLAG_READ_ONLY   = 0x0002
NBD_FLAG_HAS_FLAGS   = 0x0001


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

def virt_to_phys_segments(virt_offset, length):
    """
    Yield (phys_offset_or_None, seg_length) pairs covering the request.
    None means the region falls in a PERC metadata chunk → return zeros.
    """
    remaining = length
    pos = virt_offset
    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
        chunk_remain = CHUNK_BYTES - intra
        seg_len      = min(chunk_remain, remaining)

        if chunk_idx == 4:
            yield None, seg_len
        else:
            phys = (LV_PHYS_START
                    + group * 4 * CHUNK_BYTES
                    + chunk_idx * CHUNK_BYTES
                    + intra)
            yield phys, seg_len

        pos       += seg_len
        remaining -= seg_len


def raw_read(virt_offset, length):
    """Read virtual bytes without any patching."""
    result = bytearray(length)
    written = 0
    with open(DEV, 'rb') as f:
        for phys, seg_len in virt_to_phys_segments(virt_offset, length):
            if phys is not None:
                f.seek(phys)
                data = f.read(seg_len)
                result[written:written + len(data)] = data
                written += seg_len
            else:
                written += seg_len  # leave as zeros
    return bytes(result)


# ── Group descriptor synthesis ────────────────────────────────────────────────

def make_gdt_entry(group_num):
    """
    Build a 64-byte group descriptor for group N.

    Pattern confirmed from session forensics:
      block_bitmap_lo = 1038 + N
      inode_bitmap_lo = 1054 + N
      inode_table_lo  = 1070 + N * 512

    We set free counts to 0 and used_dirs to 0 — the kernel will recompute
    them if needed. Checksum is left zero (we cleared metadata_csum).
    """
    gd = bytearray(GDT_ENTRY)
    bb = 1038 + group_num
    ib = 1054 + group_num
    it = 1070 + group_num * 512

    struct.pack_into('<I', gd, 0,  bb)       # bg_block_bitmap_lo
    struct.pack_into('<I', gd, 4,  ib)       # bg_inode_bitmap_lo
    struct.pack_into('<I', gd, 8,  it)       # bg_inode_table_lo
    struct.pack_into('<H', gd, 12, 0)        # bg_free_blocks_count_lo
    struct.pack_into('<H', gd, 14, 0)        # bg_free_inodes_count_lo
    struct.pack_into('<H', gd, 16, 0)        # bg_used_dirs_count_lo
    struct.pack_into('<H', gd, 18, 0)        # bg_flags
    struct.pack_into('<H', gd, 30, 0)        # bg_checksum (cleared)
    return bytes(gd)


def synthesize_gdt_region(virt_offset, length):
    """
    Return synthesized GDT bytes for the virtual address range
    [virt_offset, virt_offset+length) which overlaps the GDT area.
    """
    result = bytearray(length)
    for i in range(length):
        abs_byte = virt_offset + i
        if abs_byte < GDT_START_VIRT or abs_byte >= GDT_END_VIRT:
            continue
        gdt_rel  = abs_byte - GDT_START_VIRT
        grp      = gdt_rel // GDT_ENTRY
        byte_in  = gdt_rel % GDT_ENTRY
        if grp < NUM_GROUPS:
            entry = make_gdt_entry(grp)
            result[i] = entry[byte_in]
    return bytes(result)


def overlaps(a_start, a_len, b_start, b_len):
    return a_start < b_start + b_len and b_start < a_start + a_len


# ── Superblock patching ───────────────────────────────────────────────────────

def get_patched_superblock():
    """
    Read the real superblock once, clear problematic feature bits, cache it.
    Clears:
      - has_journal      (skip journal replay)
      - metadata_csum    (stop checksum validation of bitmaps/GDT)
      - gdt_csum         (older checksum variant)
    """
    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_FEAT_INCOMPAT_OFF)[0]
        ro_compat = struct.unpack_from('<I', sb, SB_FEAT_RO_COMPAT_OFF)[0]

        incompat  &= ~INCOMPAT_HAS_JOURNAL
        ro_compat &= ~RO_COMPAT_METADATA_CSUM
        ro_compat &= ~RO_COMPAT_GDT_CSUM

        struct.pack_into('<I', sb, SB_FEAT_INCOMPAT_OFF,  incompat)
        struct.pack_into('<I', sb, SB_FEAT_RO_COMPAT_OFF, ro_compat)

        # Zero the superblock checksum (last 4 bytes) — no longer valid
        struct.pack_into('<I', sb, SB_CHECKSUM_OFF, 0)

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


# ── Main read function ────────────────────────────────────────────────────────

def read_virtual(virt_offset, length):
    """
    Read `length` bytes from virtual offset `virt_offset`, applying:
      1. Chunk translation (skip PERC metadata chunks)
      2. Superblock patching
      3. GDT synthesis for metadata-chunk regions within the GDT
    """
    # Start with raw translated data
    data = bytearray(raw_read(virt_offset, length))

    req_end = virt_offset + length

    # ── Patch 1: superblock ───────────────────────────────────────────────────
    sb_start = SB_VIRT_OFFSET
    sb_end   = SB_VIRT_OFFSET + SB_SIZE
    if overlaps(virt_offset, length, sb_start, SB_SIZE):
        patched = get_patched_superblock()
        copy_start = max(virt_offset, sb_start)
        copy_end   = min(req_end, sb_end)
        src_off    = copy_start - sb_start
        dst_off    = copy_start - virt_offset
        n          = copy_end - copy_start
        data[dst_off:dst_off + n] = patched[src_off:src_off + n]

    # ── Patch 2: GDT reconstruction ──────────────────────────────────────────
    # Only needed for regions that overlap the GDT and fall in metadata chunks
    if overlaps(virt_offset, length, GDT_START_VIRT, GDT_END_VIRT - GDT_START_VIRT):
        # Walk the metadata chunks within this request to find zero regions
        # inside the GDT and replace with synthesized data
        pos = virt_offset
        remaining = length
        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)

            if chunk_idx == 4:
                # This is a metadata chunk — was read as zeros
                # If it overlaps the GDT, synthesize
                seg_end = pos + seg_len
                if overlaps(pos, seg_len, GDT_START_VIRT,
                            GDT_END_VIRT - GDT_START_VIRT):
                    synth = synthesize_gdt_region(pos, seg_len)
                    dst_off = pos - virt_offset
                    data[dst_off:dst_off + seg_len] = synth

            pos       += seg_len
            remaining -= seg_len

    return bytes(data)


# ── NBD protocol (old-style handshake) ───────────────────────────────────────

def handle_client(conn, addr):
    print(f'[nbd] client connected from {addr}')
    try:
        # Old-style handshake: magic + cliserv_magic + size + flags + 124 pad
        conn.sendall(b'NBDMAGIC')
        conn.sendall(struct.pack('>Q', NBD_CLISERV_MAGIC))
        conn.sendall(struct.pack('>Q', VIRT_SIZE))
        conn.sendall(struct.pack('>H', NBD_FLAG_HAS_FLAGS | NBD_FLAG_READ_ONLY))
        conn.sendall(b'\x00' * 124)

        print(f'[nbd] handshake done, serving {VIRT_SIZE // (1024**3):.1f} GB')

        while True:
            hdr = b''
            while len(hdr) < 28:
                chunk = conn.recv(28 - len(hdr))
                if not chunk:
                    return
                hdr += chunk

            magic, flags, cmd, handle, offset, length = \
                struct.unpack('>IHHQQI', hdr)

            if magic != NBD_REQUEST_MAGIC:
                print(f'[nbd] bad magic 0x{magic:08x}, closing')
                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
                reply = struct.pack('>IIQ', NBD_REPLY_MAGIC, 0, handle)
                conn.sendall(reply + payload)

            elif cmd == NBD_CMD_DISC:
                print(f'[nbd] client disconnected cleanly')
                return

            else:
                # Return error for writes and other commands
                reply = struct.pack('>IIQ', NBD_REPLY_MAGIC, 1, handle)
                conn.sendall(reply)

    except (ConnectionResetError, BrokenPipeError):
        print(f'[nbd] client {addr} dropped connection')
    except Exception as e:
        print(f'[nbd] error: {e}')
    finally:
        conn.close()


def main():
    print(f'[nbd] PERC H710 recovery NBD server v2')
    print(f'[nbd] device:    {DEV}')
    print(f'[nbd] lv_start:  byte {LV_PHYS_START} (sector {LV_PHYS_START//512})')
    print(f'[nbd] virt_size: {VIRT_SIZE // (1024**3):.1f} GB')
    print(f'[nbd] features:  chunk-skip + sb-patch + gdt-synth')
    print()

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('127.0.0.1', 10809))
    server.listen(5)
    print(f'[nbd] listening on 127.0.0.1:10809')
    print(f'[nbd] connect with: nbd-client -g 127.0.0.1 10809 /dev/nbd0')
    print(f'[nbd] then mount:   mount -o ro,norecovery -t ext4 /dev/nbd0 /mnt/root')
    print()

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


if __name__ == '__main__':
    main()