#!/usr/local/bin/python3.11
"""Перевод Windows-ВМ vStack с virtio-blk на NVMe: снять запрет загрузки драйвера stornvme.

Пока ВМ работает на virtio-blk, Windows ставит драйверу NVMe (stornvme) запрет загрузки
при старте: Services\\stornvme\\StartOverride = 3. После переключения диска на NVMe
загрузчик драйвер не поднимает, и Windows падает с INACCESSIBLE_BOOT_DEVICE.
Скрипт у выключенной ВМ ставит в реестре SYSTEM на её диске запрет в 0 (и Start в 0,
если нужно) во всех ControlSet00N. Каждое изменение — 4 байта на месте: размер файлов,
метаданные NTFS и её журнал не меняются.

Запуск — на узле, где импортирован пул ВМ, от root, Python 3.11 из состава узла.
Windows-ВМ выбираются по os_type = 7 в их vm.json (/pools/<пул>/.vm_V4/<id>/vm.json), тома —
disk_* в датасете ВМ. Тип диска из vm.json не решает ничего: это тип последнего старта, и после
неудачного старта на NVMe там уже nvme. Есть ли запрет, видно только по реестру.

  win-nvme-switch.py check [VMID...]          только чтение: отчёт и вердикт
  win-nvme-switch.py patch [VMID...] [ключи]  ZFS-снимок, правка, повторная проверка

Без VMID — все Windows-ВМ пулов этого узла.
Ключи patch:
  --yes             не спрашивать подтверждение перед записью
  --no-snapshot     не делать ZFS-снимок перед правкой
  --snapshot NAME   имя снимка (по умолчанию pre-nvme-<дата>)
  --driver NAME     драйвер, которому снять запрет (по умолчанию stornvme; viostor — откат)

ВМ считается выключенной, только если сходится всё: vstack-vm показывает Offline (не
Suspended), процесса bhyve и /dev/vmm/<id> нет, и все её тома открылись монопольно
(O_EXCL). Монопольного открытия ядро не даёт, пока том открыт кем угодно — bhyve, SPDK,
ctl, — а пока скрипт держит тома, ВМ не запустится: её старт получит EBUSY.

Вердикты: READY — можно править; PATCHED — исправлено и перепроверено; OK — запрета нет;
BLOCKED — не трогаем, причина в отчёте (в том числе «ВМ работает»);
ERROR — правка начата и не закончена, разбираться руками (снимок назван в отчёте).
В конце patch — строка ids=... для win-nvme-flag.sql: ВМ в PATCHED и OK.
Код выхода: 0 — всё прошло, 2 — есть BLOCKED, 1 — есть ERROR или ошибка запуска.
"""

import datetime
import errno
import glob
import json
import os
import re
import shutil
import struct
import subprocess
import sys

CHUNK = 64 << 10
BLOCK = 4096
USA_STRIDE = 512  # шаг массива исправлений NTFS не зависит от размера сектора
HBIN0 = 4096
REG_DWORD = 4
OS_WINDOWS = 7
POOLS = "/pools"
VSTACK_VM = shutil.which("vstack-vm") or "/usr/local/sbin/vstack-vm"


class Blocked(Exception):
    """Править нельзя: причина — для администратора."""


class NotNtfs(Exception):
    """Раздел не NTFS — не ошибка, просто не системный."""


PARSE_ERRORS = (ValueError, struct.error, IndexError, UnicodeDecodeError)


# --------------------------------------------------------------------------- устройство


def who_opened(path):
    try:
        p = subprocess.run(["fstat", path], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
    except OSError:
        return "кем — fstat недоступен"
    lines = [" ".join(ln.split()[:4]) for ln in p.stdout.splitlines()[1:] if ln.strip()]
    return "; ".join(lines[:3]) if lines else "процессов нет — держит ядро (ctl/nvmf?)"


class Device:
    """Том ВМ, открытый монопольно: чтение с кэшем по 64 КиБ, запись блоками по 4 КиБ."""

    def __init__(self, path, lsect, write=False):
        self.path, self.lsect = path, lsect
        self.cache = {}
        self.fd = None
        flags = (os.O_RDWR | os.O_SYNC if write else os.O_RDONLY) | os.O_EXCL
        try:
            self.fd = os.open(path, flags)
        except OSError as e:
            if e.errno == errno.EBUSY:
                raise Blocked(f"{self.name}: том открыт ({who_opened(path)}) — ВМ не выключена "
                              "или диск отдан по сети")
            raise Blocked(f"{self.name}: открыть не удалось: {e.strerror}")
        self.size = os.lseek(self.fd, 0, os.SEEK_END)
        if self.size <= 0:
            self.close()
            raise Blocked(f"{self.name}: размер тома не определён")

    @property
    def name(self):
        return os.path.basename(self.path)

    def close(self):
        if self.fd is not None:
            os.close(self.fd)
            self.fd = None

    def drop_cache(self):
        self.cache.clear()

    def read(self, off, n):
        if n <= 0:
            return b""
        if off < 0 or off + n > self.size:
            raise ValueError(f"чтение за пределами тома: {off}+{n}")
        first, last = off // CHUNK, (off + n - 1) // CHUNK
        for c in range(first, last + 1):
            if c not in self.cache:
                self.cache[c] = os.pread(self.fd, CHUNK, c * CHUNK)
        buf = b"".join(self.cache[c] for c in range(first, last + 1))
        return buf[off - first * CHUNK:off - first * CHUNK + n]

    def read_block(self, blockno):
        data = os.pread(self.fd, BLOCK, blockno * BLOCK)
        if len(data) != BLOCK:
            raise RuntimeError(f"{self.name}: блок {blockno} прочитан не целиком")
        return data

    def write_block(self, blockno, data):
        assert len(data) == BLOCK
        if os.pwrite(self.fd, data, blockno * BLOCK) != BLOCK:
            raise RuntimeError(f"{self.name}: блок {blockno} записан не целиком")
        os.fsync(self.fd)
        self.cache.pop(blockno * BLOCK // CHUNK, None)

    @property
    def zfs_name(self):
        return self.path[len("/dev/zvol/"):] if self.path.startswith("/dev/zvol/") else None


# --------------------------------------------------------------------------- разделы


def partitions(dev):
    """[(смещение, размер)] разделов. GPT ищется на логическом секторе из vm.json, затем на 512 и 4096."""
    for sect in dict.fromkeys((dev.lsect, 512, 4096)):
        hdr = dev.read(sect, 512)
        if hdr[:8] != b"EFI PART":
            continue
        entries_lba, n, esz = struct.unpack_from("<QII", hdr, 72)
        if not (0 < n <= 4096 and 128 <= esz <= 4096):
            raise ValueError("GPT: неверный заголовок")
        tbl = dev.read(entries_lba * sect, n * esz)
        parts = []
        for i in range(n):
            e = tbl[i * esz:(i + 1) * esz]
            first, last = struct.unpack_from("<QQ", e, 32)
            if e[:16] != b"\0" * 16 and 0 < first <= last:
                parts.append((first * sect, (last - first + 1) * sect))
        return parts
    mbr = dev.read(0, 512)
    if mbr[510:512] != b"\x55\xaa":
        return []
    if mbr[3:11] == b"NTFS    ":
        return [(0, dev.size)]  # NTFS на весь диск, без таблицы разделов
    parts = []
    for i in range(4):
        e = mbr[446 + 16 * i:462 + 16 * i]
        start, cnt = struct.unpack_from("<II", e, 8)
        if e[4] not in (0x00, 0x05, 0x0F, 0x85, 0xEE) and cnt:
            parts.append((start * dev.lsect, cnt * dev.lsect))
    return parts


# --------------------------------------------------------------------------- NTFS


def decode_runs(data, off):
    """Список (lcn|None, clusters) из runlist атрибута."""
    runs, lcn = [], 0
    while off < len(data) and data[off]:
        hdr = data[off]
        ln, lo = hdr & 0x0F, hdr >> 4
        if not ln or ln > 8 or lo > 8:
            raise ValueError("runlist: неверный заголовок")
        off += 1
        count = int.from_bytes(data[off:off + ln], "little")
        off += ln
        if lo:
            lcn += int.from_bytes(data[off:off + lo], "little", signed=True)
            runs.append((lcn, count))
        else:
            runs.append((None, count))
        off += lo
    return runs


def fixup(raw, what):
    """Снять массив исправлений (update sequence) с записи NTFS на месте."""
    usa_off, usa_cnt = struct.unpack_from("<HH", raw, 4)
    if usa_cnt != len(raw) // USA_STRIDE + 1 or usa_off + 2 * usa_cnt > len(raw):
        raise ValueError(f"{what}: неверный массив исправлений")
    usn = raw[usa_off:usa_off + 2]
    for i in range(1, usa_cnt):
        pos = i * USA_STRIDE - 2
        if raw[pos:pos + 2] != usn:
            raise ValueError(f"{what}: рваная запись (fixup)")
        raw[pos:pos + 2] = raw[usa_off + 2 * i:usa_off + 2 * i + 2]


class Ntfs:
    def __init__(self, dev, base):
        self.dev, self.base = dev, base
        bs = dev.read(base, 512)
        if bs[3:11] == b"-FVE-FS-":
            raise Blocked("раздел зашифрован BitLocker")
        if bs[3:11] != b"NTFS    ":
            raise NotNtfs()
        self.sector = struct.unpack_from("<H", bs, 0x0B)[0]
        if self.sector not in (512, 1024, 2048, 4096):
            raise ValueError(f"NTFS: размер сектора {self.sector}")
        spc = bs[0x0D]
        spc = spc if spc <= 0x80 else 1 << (256 - spc)
        self.cluster = self.sector * spc
        mft_lcn = struct.unpack_from("<Q", bs, 0x30)[0]
        cpr = struct.unpack_from("<b", bs, 0x40)[0]
        self.rec_size = cpr * self.cluster if cpr > 0 else 1 << (-cpr)
        if self.rec_size not in (1024, 2048, 4096):
            raise ValueError(f"NTFS: размер записи MFT {self.rec_size}")
        self.mft_runs = [(mft_lcn, 1 << 30)]  # пока не прочитан $MFT — непрерывно
        rec0 = self.record(0)
        self.mft_runs = self.data_runs(0, rec0)

    # -- чтение по runlist
    def read_runs(self, runs, off, n):
        out = bytearray()
        vcn_off = 0
        for lcn, count in runs:
            size = count * self.cluster
            if off < vcn_off + size and n > 0:
                s = off - vcn_off
                take = min(size - s, n)
                if lcn is None:
                    out += b"\0" * take
                else:
                    out += self.dev.read(self.base + lcn * self.cluster + s, take)
                off += take
                n -= take
            vcn_off += size
            if n <= 0:
                break
        if n > 0:
            raise ValueError("чтение за концом runlist")
        return bytes(out)

    def volume_offset(self, runs, off):
        vcn_off = 0
        for lcn, count in runs:
            size = count * self.cluster
            if off < vcn_off + size:
                if lcn is None:
                    raise ValueError("разреженный участок")
                return self.base + lcn * self.cluster + (off - vcn_off)
            vcn_off += size
        raise ValueError("смещение за концом runlist")

    # -- записи MFT
    def record(self, no):
        raw = bytearray(self.read_runs(self.mft_runs, no * self.rec_size, self.rec_size))
        if raw[:4] != b"FILE":
            raise ValueError(f"запись MFT {no}: нет сигнатуры FILE")
        fixup(raw, f"запись MFT {no}")
        return bytes(raw)

    @staticmethod
    def attrs(rec):
        off = struct.unpack_from("<H", rec, 0x14)[0]
        while off + 8 <= len(rec):
            atype, alen = struct.unpack_from("<II", rec, off)
            if atype == 0xFFFFFFFF or alen == 0:
                break
            if off + alen > len(rec):
                raise ValueError("атрибут за концом записи MFT")
            a = rec[off:off + alen]
            nonres, name_len = a[8], a[9]
            name_off, flags = struct.unpack_from("<HH", a, 0x0A)
            name = a[name_off:name_off + 2 * name_len].decode("utf-16-le") if name_len else ""
            yield atype, name, nonres, flags, a
            off += alen

    def attr_pieces(self, no, rec, atype, aname=""):
        """Все куски атрибута с учётом $ATTRIBUTE_LIST: [(nonres, flags, raw)]."""
        pieces = [(nr, fl, a) for t, n, nr, fl, a in self.attrs(rec) if t == atype and n == aname]
        alist = [a for t, n, nr, fl, a in self.attrs(rec) if t == 0x20]
        if not alist:
            return pieces
        a = alist[0]
        if a[8]:
            runs = decode_runs(a, struct.unpack_from("<H", a, 0x20)[0])
            size = struct.unpack_from("<Q", a, 0x30)[0]
            data = self.read_runs(runs, 0, size)
        else:
            vlen, voff = struct.unpack_from("<IH", a, 0x10)
            data = a[voff:voff + vlen]
        refs, off = [], 0
        while off + 0x1A <= len(data):
            t, rl = struct.unpack_from("<IH", data, off)
            if rl == 0:
                break
            nl, no_ = data[off + 6], data[off + 7]
            nm = data[off + no_:off + no_ + 2 * nl].decode("utf-16-le") if nl else ""
            ref = struct.unpack_from("<Q", data, off + 0x10)[0] & 0xFFFFFFFFFFFF
            if t == atype and nm == aname and ref != no:
                refs.append(ref)
            off += rl
        for ref in dict.fromkeys(refs):
            ext = self.record(ref)
            pieces += [(nr, fl, x) for t, n, nr, fl, x in self.attrs(ext) if t == atype and n == aname]
        return pieces

    def data_runs(self, no, rec, atype=0x80, aname=""):
        pieces = [p for p in self.attr_pieces(no, rec, atype, aname) if p[0]]
        if not pieces:
            raise ValueError(f"запись MFT {no}: нет нерезидентного атрибута 0x{atype:x}")
        pieces.sort(key=lambda p: struct.unpack_from("<Q", p[2], 0x10)[0])
        runs = []
        for _, flags, a in pieces:
            if flags & 0x0001:
                raise Blocked("системный файл сжат NTFS — правка с хоста не поддерживается")
            if flags & 0x4000:
                raise Blocked("системный файл зашифрован EFS — правка с хоста невозможна")
            runs += decode_runs(a, struct.unpack_from("<H", a, 0x20)[0])
        return runs

    def data_size(self, no, rec, atype=0x80, aname=""):
        for nr, fl, a in self.attr_pieces(no, rec, atype, aname):
            if nr and struct.unpack_from("<Q", a, 0x10)[0] == 0:
                return struct.unpack_from("<Q", a, 0x30)[0]
            if not nr:
                return struct.unpack_from("<I", a, 0x10)[0]
        raise ValueError(f"запись MFT {no}: нет размера")

    def resident_value(self, rec, atype):
        for t, n, nr, fl, a in self.attrs(rec):
            if t == atype and not nr:
                vlen, voff = struct.unpack_from("<IH", a, 0x10)
                return a[voff:voff + vlen]
        return None

    # -- каталоги
    @staticmethod
    def _entries(buf, hdr_off):
        eoff, total = struct.unpack_from("<II", buf, hdr_off)
        off, end = hdr_off + eoff, min(hdr_off + total, len(buf))
        while off + 0x10 <= end:
            ref = struct.unpack_from("<Q", buf, off)[0] & 0xFFFFFFFFFFFF
            elen, klen, flags = struct.unpack_from("<HHH", buf, off + 8)
            if flags & 0x02 or elen == 0:
                break
            if klen >= 0x42:
                k = buf[off + 0x10:off + 0x10 + klen]
                nl, ns = k[0x40], k[0x41]
                yield ref, ns, k[0x42:0x42 + 2 * nl].decode("utf-16-le", errors="replace")
            off += elen

    def lookup(self, dir_no, name):
        rec = self.record(dir_no)
        want = name.lower()
        found = []
        # $INDEX_ROOT большого каталога (C:\Windows на Server 2025) лежит в записи-расширении по $ATTRIBUTE_LIST
        root = next((a[struct.unpack_from("<H", a, 0x14)[0]:][:struct.unpack_from("<I", a, 0x10)[0]]
                     for nr, fl, a in self.attr_pieces(dir_no, rec, 0x90, "$I30") if not nr), None)
        if root is None:
            raise ValueError(f"запись MFT {dir_no}: нет $INDEX_ROOT")
        idx_block = struct.unpack_from("<I", root, 8)[0]
        found += [(r, ns) for r, ns, n in self._entries(root, 0x10) if n.lower() == want]
        try:
            runs = self.data_runs(dir_no, rec, 0xA0, "$I30")
            size = self.data_size(dir_no, rec, 0xA0, "$I30")
        except ValueError:
            runs, size = [], 0
        for boff in range(0, size, idx_block):
            blk = bytearray(self.read_runs(runs, boff, idx_block))
            if blk[:4] != b"INDX":
                continue
            try:
                fixup(blk, "индексный блок")
            except ValueError:
                continue
            found += [(r, ns) for r, ns, n in self._entries(bytes(blk), 0x18) if n.lower() == want]
        # Индексные блоки могут хранить устаревшие записи удалённых файлов: берём только
        # запись MFT, которая занята и чьё имя действительно лежит в этом каталоге.
        for ref, ns in sorted(found, key=lambda x: x[1] == 2):
            if self._is_child(ref, dir_no, want):
                return ref
        return None

    def _is_child(self, no, parent, want):
        try:
            rec = self.record(no)
        except ValueError:
            return False
        if not struct.unpack_from("<H", rec, 0x16)[0] & 0x01:
            return False
        for nr, fl, a in self.attr_pieces(no, rec, 0x30):
            if not nr:
                vlen, voff = struct.unpack_from("<IH", a, 0x10)
                v = a[voff:voff + vlen]
                par = struct.unpack_from("<Q", v, 0)[0] & 0xFFFFFFFFFFFF
                nm = v[0x42:0x42 + 2 * v[0x40]].decode("utf-16-le", errors="replace")
                if par == parent and nm.lower() == want:
                    return True
        return False

    def path(self, path):
        no = 5
        for part in path.split("/"):
            no = self.lookup(no, part)
            if no is None:
                return None
        return no

    def volume_dirty(self):
        info = self.resident_value(self.record(3), 0x70)
        return bool(info and struct.unpack_from("<H", info, 10)[0] & 0x0001)

    def log_clean(self):
        """Журнал NTFS закрыт чисто: том отмонтирован штатно, Windows повторять в нём нечего.

        Та же проверка, что у ntfs-3g: в рабочей странице рестарта нет активных клиентов
        журнала либо стоит флаг RESTART_VOLUME_IS_CLEAN.
        """
        f = NtfsFile(self, 2)
        head = f.read(0, 65536)
        if head and head == b"\xff" * len(head):
            return True  # журнал пуст: его сбросили chkdsk или ntfs-3g, повторять нечего
        pages, off = [], 0
        for _ in range(2):
            hdr = f.read(off, 0x20)
            psize = struct.unpack_from("<I", hdr, 0x10)[0] if hdr[:4] == b"RSTR" else 4096
            if psize not in (4096, 8192, 16384, 32768, 65536):
                psize = 4096
            raw = bytearray(f.read(off, psize))
            if raw[:4] == b"RSTR":
                try:
                    fixup(raw, "$LogFile")
                    ra = struct.unpack_from("<H", raw, 0x18)[0]
                    lsn, _, _, in_use, flags = struct.unpack_from("<QHHHH", raw, ra)
                    pages.append((lsn, in_use, flags))
                except PARSE_ERRORS:
                    pass
            off = psize
        if not pages:
            raise Blocked("журнал NTFS ($LogFile) не читается — состояние тома неизвестно")
        _, in_use, flags = max(pages)
        return in_use == 0xFFFF or bool(flags & 0x0002)


class NtfsFile:
    def __init__(self, fs, no):
        self.fs, self.no = fs, no
        rec = fs.record(no)
        self.runs = fs.data_runs(no, rec)
        self.size = fs.data_size(no, rec)

    def read(self, off, n):
        n = max(0, min(n, self.size - off))
        return self.fs.read_runs(self.runs, off, n)

    def volume_offset(self, off):
        return self.fs.volume_offset(self.runs, off)


# --------------------------------------------------------------------------- реестр


def hive_checksum(hdr):
    c = 0
    for i in range(0, 508, 4):
        c ^= struct.unpack_from("<I", hdr, i)[0]
    return 0xFFFFFFFE if c == 0xFFFFFFFF else (1 if c == 0 else c)


class Hive:
    def __init__(self, f):
        self.f = f
        hdr = f.read(0, 512)
        if hdr[:4] != b"regf":
            raise Blocked("файл реестра повреждён (нет сигнатуры regf)")
        self.seq1, self.seq2 = struct.unpack_from("<II", hdr, 4)
        self.root_off = struct.unpack_from("<I", hdr, 0x24)[0]
        self.bins_size = struct.unpack_from("<I", hdr, 0x28)[0]
        self.checksum_ok = hive_checksum(hdr) == struct.unpack_from("<I", hdr, 0x1FC)[0]

    @property
    def clean(self):
        return self.seq1 == self.seq2 and self.checksum_ok

    def cell(self, off):
        if off + 4 > self.bins_size:
            raise ValueError(f"ячейка реестра 0x{off:x} за концом данных")
        size = abs(struct.unpack_from("<i", self.f.read(HBIN0 + off, 4))[0])
        if size < 8 or off + size > self.bins_size:
            raise ValueError(f"ячейка реестра 0x{off:x}: неверный размер {size}")
        return self.f.read(HBIN0 + off + 4, size - 4)

    def _name(self, cell, nlen_off, name_off, comp):
        nlen = struct.unpack_from("<H", cell, nlen_off)[0]
        raw = cell[name_off:name_off + nlen]
        return raw.decode("latin-1") if comp else raw.decode("utf-16-le", errors="replace")

    def key_name(self, nk):
        return self._name(nk, 0x48, 0x4C, struct.unpack_from("<H", nk, 2)[0] & 0x20)

    def subkeys(self, nk):
        cnt, lst = struct.unpack_from("<I", nk, 0x14)[0], struct.unpack_from("<I", nk, 0x1C)[0]
        if not cnt or lst == 0xFFFFFFFF:
            return
        stack = [lst]
        while stack:
            c = self.cell(stack.pop())
            sig, n = c[:2], struct.unpack_from("<H", c, 2)[0]
            if sig in (b"lf", b"lh"):
                offs = [struct.unpack_from("<I", c, 4 + 8 * i)[0] for i in range(n)]
            elif sig in (b"li", b"ri"):
                offs = [struct.unpack_from("<I", c, 4 + 4 * i)[0] for i in range(n)]
            else:
                raise ValueError(f"неизвестный список подключей {sig!r}")
            if sig == b"ri":
                stack += offs
                continue
            for o in offs:
                nk = self.cell(o)
                if nk[:2] == b"nk":
                    yield o, nk

    def open(self, path):
        nk = self.cell(self.root_off)
        for part in path.split("\\"):
            want = part.lower()
            for _, sub in self.subkeys(nk):
                if self.key_name(sub).lower() == want:
                    nk = sub
                    break
            else:
                return None
        return nk

    def values(self, nk):
        cnt, lst = struct.unpack_from("<II", nk, 0x24)
        if not cnt or lst == 0xFFFFFFFF:
            return
        offs = struct.unpack_from(f"<{cnt}I", self.cell(lst), 0)
        for o in offs:
            vk = self.cell(o)
            if vk[:2] != b"vk":
                continue
            name = self._name(vk, 2, 0x14, struct.unpack_from("<H", vk, 0x10)[0] & 1) or "(Default)"
            dlen, data, vtype = struct.unpack_from("<III", vk, 4)
            yield o, name, vtype, dlen, data

    def value(self, nk, name):
        for v in self.values(nk) or ():
            if v[1].lower() == name.lower():
                return v
        return None

    def string(self, nk, name):
        v = self.value(nk, name) if nk is not None else None
        if not v:
            return None
        _, _, vtype, dlen, data = v
        size = dlen & 0x7FFFFFFF
        if dlen & 0x80000000:
            raw = struct.pack("<I", data)[:size]
        elif size > 16344:
            return None
        else:
            raw = self.cell(data)[:size]
        if vtype in (1, 2):
            return raw.decode("utf-16-le", errors="replace").rstrip("\0")
        if vtype == REG_DWORD and size == 4:
            return str(struct.unpack("<I", raw)[0])
        return None


# --------------------------------------------------------------------------- анализ тома


class Change:
    def __init__(self, where, hive_off, old):
        self.where, self.hive_off, self.old = where, hive_off, old

    def key(self):
        return self.where, self.hive_off, self.old


def inline_dword(vtype, dlen):
    return vtype == REG_DWORD and bool(dlen & 0x80000000) and (dlen & 0x7FFFFFFF) == 4


def analyse_windows(fs, hive_no, driver, res):
    """Разбор системного раздела Windows; Blocked — править нельзя."""
    info = res["info"]
    sw_no = fs.path("Windows/System32/config/SOFTWARE")
    if sw_no is not None:
        try:
            sw = Hive(NtfsFile(fs, sw_no))
            cv = sw.open("Microsoft\\Windows NT\\CurrentVersion")
            prod, build = sw.string(cv, "ProductName"), sw.string(cv, "CurrentBuildNumber")
            info.append(f"Windows: {prod or '?'}, сборка {build or '?'}")
        except (Blocked,) + PARSE_ERRORS:
            info.append("Windows: версию прочитать не удалось")
    if fs.volume_dirty():
        raise Blocked("том NTFS помечен грязным — запустите ВМ как есть (virtio-blk), "
                      "выключите из Windows через «Завершение работы» и повторите")
    if not fs.log_clean():
        raise Blocked("журнал NTFS не закрыт: ВМ выключили не штатно или Windows в гибернации/"
                      "быстром запуске — запустите ВМ как есть, powercfg /h off, "
                      "«Завершение работы», повторите")
    hib = fs.path("hiberfil.sys")
    if hib is not None:
        try:
            sig = NtfsFile(fs, hib).read(0, 4)
        except PARSE_ERRORS:
            sig = b""
        if sig.lower() in (b"hibr", b"wake", b"rstr"):
            raise Blocked("в hiberfil.sys сохранено ядро (гибернация или быстрый запуск) — "
                          "в Windows: powercfg /h off, выключить, повторить")
    f = NtfsFile(fs, hive_no)
    res["file"] = f
    hive = Hive(f)
    if not hive.clean:
        raise Blocked(f"реестр SYSTEM не чистый (seq {hive.seq1}/{hive.seq2}, "
                      f"контрольная сумма {'ok' if hive.checksum_ok else 'неверна'}) — "
                      "запустите ВМ как есть и выключите её из Windows, затем повторите")
    sets = []
    for _, nk in hive.subkeys(hive.cell(hive.root_off)):
        nm = hive.key_name(nk)
        if nm.lower().startswith("controlset") and nm[10:].isdigit():
            sets.append(nm)
    sel = hive.open("Select")
    cur = hive.string(sel, "Current") if sel is not None else None
    info.append(f"наборы настроек: {', '.join(sorted(sets))}; загрузочный — "
                f"ControlSet{int(cur or 0):03d}")
    seen = False
    for cs in sorted(sets):
        svc = hive.open(f"{cs}\\Services\\{driver}")
        if svc is None:
            info.append(f"{cs}: службы {driver} нет")
            continue
        seen = True
        st = hive.value(svc, "Start")
        if st is None:
            raise Blocked(f"{cs}\\Services\\{driver}: нет значения Start")
        o, _, vtype, dlen, data = st
        if not inline_dword(vtype, dlen):
            raise Blocked(f"{cs}\\Services\\{driver}\\Start: неожиданный тип {vtype}")
        line = f"{cs}: Start={data}"
        if data != 0:
            res["changes"].append(Change(f"{cs}\\Services\\{driver}\\Start", HBIN0 + o + 4 + 8, data))
        so_key = None
        for _, sub in hive.subkeys(svc):
            if hive.key_name(sub).lower() == "startoverride":
                so_key = sub
        if so_key is None:
            line += ", StartOverride нет"
        else:
            vals = []
            for vo, vname, vtype, dlen, data in hive.values(so_key) or ():
                if not inline_dword(vtype, dlen):
                    raise Blocked(f"{cs}\\...\\StartOverride\\{vname}: неожиданный тип {vtype}")
                vals.append(f"{vname}={data}")
                # любое ненулевое значение не даёт драйверу стартовать при загрузке
                if data != 0:
                    res["changes"].append(Change(f"{cs}\\Services\\{driver}\\StartOverride\\{vname}",
                                                 HBIN0 + vo + 4 + 8, data))
            line += f", StartOverride: {', '.join(vals) or 'пусто'}"
        info.append(line)
    if not seen:
        raise Blocked(f"в системе нет драйвера {driver} — Windows без встроенной поддержки NVMe "
                      "(старше 8.1 / Server 2012 R2), переводить нельзя")


def analyse_disk(dev, driver):
    """dict: windows, info[], changes[], file, notes[]. Blocked — системный раздел, править нельзя."""
    res = {"windows": False, "info": [], "changes": [], "file": None, "notes": []}
    parts = partitions(dev)
    if not parts:
        res["notes"].append(f"{dev.name}: таблицы разделов нет")
        return res
    for off, size in parts:
        where = f"{dev.name}: раздел {off // 2**20} МиБ от начала"
        try:
            fs = Ntfs(dev, off)
            hive_no = fs.path("Windows/System32/config/SYSTEM")
        except NotNtfs:
            continue
        except Blocked as e:
            res["notes"].append(f"{where}: {e}")
            continue
        except PARSE_ERRORS as e:
            res["notes"].append(f"{where}: NTFS не разбирается ({e})")
            continue
        if hive_no is None:
            continue
        if res["windows"]:
            raise Blocked(f"{dev.name}: на диске два раздела с Windows — какой загрузочный, не ясно")
        res["windows"] = True
        res["info"].append(f"раздел NTFS {off // 2**20} МиБ от начала, {size // 2**30} ГиБ")
        try:
            analyse_windows(fs, hive_no, driver, res)
        except PARSE_ERRORS as e:
            raise Blocked(f"{dev.name}: системный раздел не разбирается ({e})")
    return res


# --------------------------------------------------------------------------- правка


def apply_changes(dev, res, log, state):
    f = res["file"]
    for ch in res["changes"]:
        voff = f.volume_offset(ch.hive_off)
        blockno, pos = divmod(voff, BLOCK)
        if pos > BLOCK - 4:
            raise RuntimeError(f"{ch.where}: значение на границе блока — не пишу")
        blk = bytearray(dev.read_block(blockno))
        cur = struct.unpack_from("<I", blk, pos)[0]
        via_file = struct.unpack("<I", f.read(ch.hive_off, 4))[0]
        if cur != ch.old or via_file != ch.old:
            raise RuntimeError(f"{ch.where}: на томе {cur}, через файл {via_file}, ожидал {ch.old} — не пишу")
        struct.pack_into("<I", blk, pos, 0)
        state["written"] = True
        dev.write_block(blockno, bytes(blk))
        after = struct.unpack_from("<I", dev.read_block(blockno), pos)[0]
        if after != 0:
            raise RuntimeError(f"{ch.where}: после записи читается {after}")
        log(f"записано: {dev.name} {ch.where} {ch.old} -> 0 (смещение на томе {voff})")


# --------------------------------------------------------------------------- ВМ


class Vm:
    def __init__(self, vmid, json_path):
        self.id, self.json_path = vmid, json_path
        self.name, self.pool, self.os_type = "?", json_path.split("/")[2], None
        with open(json_path) as fh:
            j = json.load(fh)
        if int(j.get("id", -1)) != vmid:
            raise Blocked(f"{json_path}: id в файле {j.get('id')} не совпадает с каталогом")
        self.name = j.get("name") or "?"
        self.pool = j.get("pool_name") or self.pool
        self.os_type = j.get("os_type")
        self.root_ds = j.get("root_dataset_name") or f"{self.pool}/.vm_V4/{vmid:09d}"
        disks = j.get("disks") or []
        # devtype в vm.json — тип диска последнего старта, а не то, грузится ли с него Windows:
        # после неудачного старта на NVMe там уже nvme. Судить можно только по реестру.
        self.devtypes = sorted({d.get("devtype") or "virtio-blk" for d in disks})
        self.lsect = {d["dev"]: int((d.get("sector_size") or {}).get("logical") or 512)
                      for d in disks if d.get("dev")}

    @property
    def label(self):
        return f"{self.id} ({self.name})"

    def volumes(self):
        """[(путь, логический сектор)] — тома disk_* в датасете ВМ на этом узле.

        Не из vm.json: после старта на чужом узле там вместо тома адрес NVMe-oF.
        """
        p = subprocess.run(["zfs", "list", "-H", "-o", "name", "-t", "volume", "-r", "-d", "1", self.root_ds],
                           stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if p.returncode != 0:
            raise Blocked(f"датасета {self.root_ds} на узле нет: {p.stderr.strip()}")
        vols = sorted(n for n in p.stdout.split() if n.rsplit("/", 1)[-1].startswith("disk_"))
        if not vols:
            raise Blocked(f"в {self.root_ds} нет томов disk_*")
        return [(f"/dev/zvol/{v}", self.lsect.get(f"/dev/zvol/{v}", 512)) for v in vols]


def find_vms(ids):
    """[(vmid, Vm | Blocked)] — ВМ из vm.json на пулах узла."""
    found = {}
    for path in sorted(glob.glob(f"{POOLS}/*/.vm_V4/*/vm.json")):
        d = os.path.basename(os.path.dirname(path))
        if d.isdigit():
            found.setdefault(int(d), []).append(path)
    out = []
    for vmid in (ids or sorted(found)):
        paths = found.get(vmid)
        try:
            if not paths:
                raise Blocked("vm.json на этом узле нет — пул импортирован на другом узле?")
            if len(paths) > 1:
                raise Blocked(f"vm.json в нескольких пулах: {', '.join(paths)}")
            vm = Vm(vmid, paths[0])
        except (OSError, json.JSONDecodeError, TypeError, ValueError) as e:
            out.append((vmid, Blocked(f"vm.json не читается: {e}")))
            continue
        except Blocked as e:
            out.append((vmid, e))
            continue
        if vm.os_type != OS_WINDOWS:
            if ids:
                out.append((vmid, Blocked(f"не Windows: os_type={vm.os_type} в vm.json")))
            continue
        out.append((vmid, vm))
    return out


def vm_processes(vmid):
    p = subprocess.run(["ps", "-axww", "-o", "pid=,command="], stdout=subprocess.PIPE, text=True, check=True)
    pat = re.compile(rf"^\s*(\d+)\s+(\S*bhyve: 0*{vmid})(\s|$)")
    return [m.group(1) + " " + m.group(2) for m in map(pat.match, p.stdout.splitlines()) if m]


def ensure_no_process(vm):
    procs = vm_processes(vm.id)
    if procs:
        raise Blocked(f"ВМ работает: {'; '.join(procs)}")
    for name in dict.fromkeys((str(vm.id), f"{vm.id:09d}")):
        if os.path.exists(f"/dev/vmm/{name}"):
            raise Blocked(f"есть /dev/vmm/{name} — ВМ не выключена или не освобождена")


def ensure_offline(vm):
    p = subprocess.run([VSTACK_VM, "vm", "status", f"--id={vm.id}"],
                       stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if p.returncode != 0:
        raise Blocked(f"vstack-vm vm status: {p.stderr.strip() or 'код ' + str(p.returncode)}")
    rows = [ln.split() for ln in p.stdout.splitlines()]
    row = next((r for r in rows if len(r) >= 3 and r[0].lstrip("0") == str(vm.id).lstrip("0")), None)
    status = row[2] if row else "?"
    if status == "Suspended":
        raise Blocked("ВМ приостановлена (Suspended): возобновить, выключить штатно, повторить")
    if status != "Offline":
        raise Blocked(f"vstack-vm: состояние {status}, а не Offline")
    ensure_no_process(vm)


def inspect(vm, driver, write):
    """Доказать, что ВМ выключена, монопольно захватить её тома и разобрать их.

    Возвращает (devs, results) — тома остаются захваченными, закрыть их должен вызывающий.
    """
    vols = vm.volumes()
    ensure_offline(vm)
    devs = []
    try:
        for path, lsect in vols:
            devs.append(Device(path, lsect, write))
        ensure_no_process(vm)  # ещё раз, уже под захватом
        results = [analyse_disk(dev, driver) for dev in devs]
        if not any(r["windows"] for r in results):
            notes = [n for r in results for n in r["notes"]]
            raise Blocked("системный раздел Windows на дисках ВМ не найден"
                          + (": " + "; ".join(notes) if notes else ""))
    except BaseException:
        for dev in devs:
            dev.close()
        raise
    return devs, results


def signature(devs, results):
    return [(dev.path, ch.key()) for dev, r in zip(devs, results) for ch in r["changes"]]


def report(devs, results, log):
    for dev, r in zip(devs, results):
        for n in r["notes"]:
            log(n)
        for ln in r["info"]:
            log(f"{dev.name}: {ln}")
    for dev, r in zip(devs, results):
        for ch in r["changes"]:
            log(f"править: {dev.name} {ch.where} = {ch.old} -> 0")


def check_vm(vm, driver):
    """(вердикт, причина, сигнатура правок)."""
    log = lambda s: print(f"   {s}")
    devs, results = inspect(vm, driver, write=False)
    try:
        log("выключена: vstack-vm Offline, bhyve и /dev/vmm нет, тома захвачены монопольно")
        log(f"последний старт в vm.json: {', '.join(vm.devtypes) or '?'}")
        report(devs, results, log)
        sig = signature(devs, results)
        return ("READY" if sig else "OK"), "", sig
    finally:
        for dev in devs:
            dev.close()


def patch_vm(vm, driver, snap, expected):
    """(вердикт, причина). Тома держатся монопольно от проверки до повторного разбора."""
    log = lambda s: print(f"   {s}")
    devs, results = inspect(vm, driver, write=True)
    state, snaps = {"written": False}, []
    try:
        if signature(devs, results) != expected:
            raise Blocked("состояние дисков изменилось после проверки — повторите check")
        todo = [(dev, r) for dev, r in zip(devs, results) if r["changes"]]
        if snap:
            for dev, _ in todo:
                name = f"{dev.zfs_name}@{snap}"
                p = subprocess.run(["zfs", "snapshot", name], stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE, text=True)
                if p.returncode != 0:
                    raise Blocked(f"снимок {name} не сделан: {p.stderr.strip()}")
                snaps.append(name)
                log(f"снимок: {name}")
        for dev, r in todo:
            apply_changes(dev, r, log, state)
        for dev in devs:
            dev.drop_cache()
        again = [analyse_disk(dev, driver) for dev in devs]
        left = [f"{dev.name} {ch.where}={ch.old}" for dev, r in zip(devs, again) for ch in r["changes"]]
        if left:
            raise RuntimeError(f"после правки остались: {', '.join(left)}")
        log("повторная проверка: запретов нет")
        return "PATCHED", ""
    except Exception as e:
        if state["written"]:
            return "ERROR", f"{e}; снимки: {', '.join(snaps) or 'нет'}"
        if isinstance(e, Blocked):
            raise
        raise Blocked(f"{e} (на диск ничего не записано; снимки: {', '.join(snaps) or 'нет'})")
    finally:
        for dev in devs:
            dev.close()


# --------------------------------------------------------------------------- запуск


def main(argv):
    if not argv or argv[0] not in ("check", "patch"):
        print(__doc__)
        return 1
    mode, rest = argv[0], argv[1:]
    driver, snap, do_snap, yes, ids = "stornvme", None, True, False, []
    i = 0
    while i < len(rest):
        a = rest[i]
        if a in ("--driver", "--snapshot"):
            if i + 1 >= len(rest):
                print(f"{a}: нет значения", file=sys.stderr)
                return 1
            if a == "--driver":
                driver = rest[i + 1]
            else:
                snap = rest[i + 1]
            i += 1
        elif a == "--no-snapshot":
            do_snap = False
        elif a == "--yes":
            yes = True
        elif a.isdigit():
            ids.append(int(a))
        else:
            print(f"непонятный аргумент: {a}", file=sys.stderr)
            return 1
        i += 1
    if os.geteuid() != 0:
        print("нужен root", file=sys.stderr)
        return 1
    snap = (snap or "pre-nvme-" + datetime.datetime.now().strftime("%Y%m%d-%H%M")) if do_snap else None

    vms = find_vms(list(dict.fromkeys(ids)))
    if not vms:
        print("Windows-ВМ на пулах этого узла нет.")
        return 0

    results = {}  # vmid -> (vm, вердикт, причина)
    ready = []
    for vmid, vm in vms:
        if isinstance(vm, Blocked):
            print(f"== ВМ {vmid}")
            print(f"   ВЕРДИКТ: BLOCKED — {vm}")
            results[vmid] = (None, "BLOCKED", str(vm))
            continue
        print(f"== ВМ {vm.label}, пул {vm.pool}")
        try:
            verdict, reason, sig = check_vm(vm, driver)
        except Blocked as e:
            verdict, reason, sig = "BLOCKED", str(e), None
        print(f"   ВЕРДИКТ: {verdict}{' — ' + reason if reason else ''}")
        results[vmid] = (vm, verdict, reason)
        if verdict == "READY":
            ready.append((vm, sig))

    if mode == "patch" and ready:
        print()
        if not yes:
            if not sys.stdin.isatty():
                print("Без терминала правка только с --yes.", file=sys.stderr)
                return 1
            ans = input(f"Править {len(ready)} ВМ ({', '.join(str(v.id) for v, _ in ready)}), "
                        f"драйвер {driver}, снимок {'@' + snap if snap else 'не делать'}? [yes/N] ")
            if ans.strip().lower() != "yes":
                print("Отменено.")
                return 1
        for vm, sig in ready:
            print(f"== ВМ {vm.label}: правка")
            try:
                verdict, reason = patch_vm(vm, driver, snap, sig)
            except Blocked as e:
                verdict, reason = "BLOCKED", str(e)
            print(f"   ИТОГ: {verdict}{' — ' + reason if reason else ''}")
            results[vm.id] = (vm, verdict, reason)

    def group(*verdicts):
        return [(vmid, vm, r) for vmid, (vm, v, r) in sorted(results.items()) if v in verdicts]

    def names(rows, why=False):
        return "; ".join(f"{vm.label if vm else vmid}" + (f" — {r}" if why and r else "")
                         for vmid, vm, r in rows) or "—"

    print("\n== Итог")
    if mode == "patch":
        print(f"   драйвер {driver} включён: {names(group('PATCHED'))}")
        print(f"   запрета не было: {names(group('OK'))}")
    else:
        print(f"   к правке: {names(group('READY'))}")
        print(f"   запрета нет: {names(group('OK'))}")
    if group("BLOCKED"):
        print(f"   не тронуты: {names(group('BLOCKED'), True)}")
    if group("ERROR"):
        print(f"   ОШИБКА ПРАВКИ: {names(group('ERROR'), True)}")
    if mode == "patch" and driver == "stornvme":
        ok = [str(vmid) for vmid, _, _ in group("PATCHED", "OK")]
        print(f"ids={','.join(ok)}" if ok else "ids= (переводить на NVMe нечего)")

    if group("ERROR"):
        return 1
    return 2 if group("BLOCKED") else 0


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv[1:]))
    except KeyboardInterrupt:
        sys.exit(1)