#!/usr/bin/env python3
"""Restore a NAS Reborn backup without the app.

This script is the reference reader for the NAS Reborn vault format and
doubles as its specification: every constant a reader needs is written
down below. It needs the vault as a folder on a computer, the recovery
phrase, Python 3 and the "cryptography" package:

    python3 -m pip install cryptography

    python3 nas-reborn-restore.py list     VAULT
    python3 nas-reborn-restore.py restore  VAULT OUTPUT [--snapshot N] [--path PREFIX]
    python3 nas-reborn-restore.py verify   VAULT [--snapshot N]
    python3 nas-reborn-restore.py rescue   VAULT OUTPUT
    python3 nas-reborn-restore.py self-test

VAULT is the folder holding "nas-reborn-backup", or that folder itself.
The recovery phrase is asked for interactively, or read from a file
given with --phrase-file, or from standard input when that is not a
terminal. Never pass it as an argument: shells keep those in their
history.

"restore" writes the newest snapshot, or the one named with --snapshot,
into OUTPUT, as one folder per backed-up folder. --path restricts it to
one folder or file inside the snapshot. "verify" decrypts everything a
snapshot references and writes nothing. "rescue" is for a vault whose
snapshot indexes are all lost or unreadable: it rebuilds a best guess
from the encrypted headers of the blobs and restores that.

Without the recovery phrase nothing here can help. The copy of the key
kept in iCloud Keychain can only be used by another Apple device signed
into the same account, from inside the app.


THE VAULT FORMAT, VERSION 1
===========================

Layout
------

Everything lives under one folder, which is one key prefix on an S3
bucket and one plain directory on a NAS Reborn host, where object keys
map onto file paths one to one:

    nas-reborn-backup/README.txt             plain text, says what this is
    nas-reborn-backup/nas-reborn-restore.py  this script, plain text
    nas-reborn-backup/claim                  plain text, names the writing device
    nas-reborn-backup/m-0000000042           a sealed snapshot index, one per
                                             generation, zero-padded to ten
                                             digits so a plain sort is history
    nas-reborn-backup/b/<name>               sealed file blobs; <name> is 32
                                             lowercase Base32 characters

Restoring needs only the m- and b/ objects and the master key. Nothing
is kept in S3 metadata or in extended attributes, so a copy of the
files is a copy of the backup.

Keys
----

The master key is 32 random bytes, made on the device that backs up and
never sent to the vault. Three working keys derive from it with
HKDF-SHA256, empty salt, 32 bytes out, and these info strings:

    "nas-vault-v1 name"        the name key, HMACs paths into object names
    "nas-vault-v1 header"      the header key, seals blob headers
    "nas-vault-v1 manifest"    the manifest key, seals snapshot indexes

Recovery phrase
---------------

The master key followed by the first two bytes of its SHA-256, 34 bytes
in all, encoded as lowercase Base32 and written as 55 characters in
groups of five separated by "-". Readers ignore case, spaces and
separators, and check the checksum so a mistyped phrase never derives
wrong keys silently.

Base32
------

The RFC 4648 alphabet in lowercase, "abcdefghijklmnopqrstuvwxyz234567",
without padding characters; the final partial group is filled with zero
bits. Used for object names and for the phrase.

Sealed boxes
------------

Every encrypted thing is AES-256-GCM in one shape,

    nonce (12 bytes) || ciphertext || tag (16 bytes)

with associated data (AAD) as given for each use. The 28 bytes of nonce
and tag are the seal overhead. A fresh random nonce is used for every
box.

Snapshot index (manifest)
-------------------------

One sealed box under the manifest key, with AAD = the generation number
as 8 bytes big-endian, so a host can never answer a request for one
generation with another. The plaintext is JSON:

    {"version": 1,
     "generation": 42,
     "created": <date>,
     "carriedForwardFolders": ["<folder>", ...],
     "entries": [
        {"path": "<folder>/<file path>",
         "objectName": "<32 Base32 characters>",
         "size": <bytes>,
         "mtime": <date>,
         "contentSHA256": "<base64>"},
        ...]}

Each index is a complete listing of one snapshot, never a diff. The
retained indexes are the snapshots. "carriedForwardFolders" names
top-level folders that were unreachable when the snapshot was made,
whose entries were copied from the previous snapshot unchanged.

Dates and binary in JSON: a date is a number of seconds, with fraction,
since 2001-01-01 00:00:00 UTC (Apple's reference date; add 978307200
for a Unix time). Binary is standard base64 with padding.

Object names
------------

Base32 of the first 20 bytes of

    HMAC-SHA256(name key, path || 0x00 || SHA-256 of the plaintext)

giving 32 characters. A name commits to path and content both, so an
unchanged file keeps its name across snapshots and a changed file lands
under a new name instead of overwriting the old version.

Blob
----

A file object is laid out as

    u32 big-endian    length of the header box
    header box        sealed under the header key,
                      AAD = object name (UTF-8) || "h"
    segment box 0     sealed under the content key from the header,
    segment box 1     AAD = object name (UTF-8) ||
    ...                     segment index as 8 bytes big-endian ||
    segment box n-1         one byte, 1 for the final segment and else 0

with this JSON as the header plaintext:

    {"version": 1,
     "path": "<folder>/<file path>",
     "size": <bytes>,
     "mtime": <date>,
     "contentSHA256": "<base64>",
     "contentKey": "<base64, 32 bytes>"}

The content is the plaintext cut into segments of 2 MiB (2097152
bytes), each sealed under the header's content key; the final segment
holds the remainder and is flagged in its AAD. An empty file has
exactly one final segment with no plaintext. A blob for a file of
`size` bytes therefore has max(1, ceil(size / 2 MiB)) segments and is
exactly

    4 + header box length + size + 28 * segment count

bytes long, which a reader checks before opening anything: shorter is
truncation, longer is not a blob. Every tag is verified before its
plaintext is used, and the header's path and hash must agree with the
index entry that led to the blob.

Rescue
------

Because every blob's header carries its path, size, mtime and hash, a
vault with no readable index can still be read: open every blob's
header, and keep per path the one with the latest mtime. There is then
no record of which version was current at the end, so this is a best
guess and is presented as one.
"""

import argparse
import base64
import datetime
import getpass
import hashlib
import hmac
import json
import os
import re
import struct
import sys

try:
    from cryptography.exceptions import InvalidTag
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
except ImportError:
    sys.exit("this script needs the cryptography package:  python3 -m pip install cryptography")

FORMAT_VERSION = 1
ROOT = "nas-reborn-backup"
SEGMENT_SIZE = 2 * 1024 * 1024
NONCE_SIZE = 12
TAG_SIZE = 16
SEAL_OVERHEAD = NONCE_SIZE + TAG_SIZE
MAX_HEADER_BOX = SEAL_OVERHEAD + 16384      # a header is small; anything larger is not one
REFERENCE_DATE = 978307200                  # 2001-01-01T00:00:00Z as a Unix time
B32_ALPHABET = "abcdefghijklmnopqrstuvwxyz234567"
MANIFEST_NAME = re.compile(r"^m-(\d{10})$")
BLOB_NAME = re.compile(r"^[a-z2-7]{32}$")


class FormatError(Exception):
    """the bytes are not what the format says they should be"""


# MARK: primitives

def b32encode(data):
    out, buffer, bits = [], 0, 0
    for byte in data:
        buffer = ((buffer << 8) | byte) & 0xFFFF
        bits += 8
        while bits >= 5:
            bits -= 5
            out.append(B32_ALPHABET[(buffer >> bits) & 0x1F])
    if bits:
        out.append(B32_ALPHABET[(buffer << (5 - bits)) & 0x1F])
    return "".join(out)


def b32decode(text):
    out, buffer, bits = bytearray(), 0, 0
    for char in text:
        value = B32_ALPHABET.find(char)
        if value < 0:
            raise ValueError("not a Base32 character: %r" % char)
        buffer = ((buffer << 5) | value) & 0xFFFF
        bits += 5
        if bits >= 8:
            bits -= 8
            out.append((buffer >> bits) & 0xFF)
    # a valid unpadded encoding leaves at most 4 bits, and they are zero
    if bits > 4 or buffer & ((1 << bits) - 1):
        raise ValueError("malformed Base32")
    return bytes(out)


def open_box(key, box, aad):
    if len(box) < SEAL_OVERHEAD:
        raise FormatError("sealed box too short")
    try:
        return AESGCM(key).decrypt(box[:NONCE_SIZE], box[NONCE_SIZE:], aad)
    except InvalidTag:
        raise FormatError("authentication failed: wrong key, or the object was altered")


def unix_time(apple_date):
    return apple_date + REFERENCE_DATE


def format_date(apple_date):
    return datetime.datetime.fromtimestamp(unix_time(apple_date)).strftime("%Y-%m-%d %H:%M")


def human(size):
    for unit in ("bytes", "KB", "MB", "GB", "TB"):
        if size < 1000 or unit == "TB":
            return ("%d %s" if unit == "bytes" else "%.1f %s") % (size, unit)
        size /= 1000.0


class Keys:
    def __init__(self, master):
        if len(master) != 32:
            raise ValueError("the master key is 32 bytes")
        self.master = master
        self.name = self._derive(b"name")
        self.header = self._derive(b"header")
        self.manifest = self._derive(b"manifest")

    def _derive(self, info):
        return HKDF(algorithm=hashes.SHA256(), length=32, salt=None,
                    info=b"nas-vault-v1 " + info).derive(self.master)

    @classmethod
    def from_phrase(cls, phrase):
        cleaned = "".join(c for c in phrase.lower() if c not in "- _\n\r\t.")
        try:
            payload = b32decode(cleaned)
        except ValueError:
            raise ValueError("the phrase holds a character that cannot be part of one")
        if len(payload) != 34:
            raise ValueError("a recovery phrase is 55 characters, this one decodes to %d bytes" % len(payload))
        master, checksum = payload[:32], payload[32:]
        if hashlib.sha256(master).digest()[:2] != checksum:
            raise ValueError("the phrase fails its checksum: a character is wrong or missing")
        return cls(master)

    def phrase(self):
        flat = b32encode(self.master + hashlib.sha256(self.master).digest()[:2])
        return "-".join(flat[i:i + 5] for i in range(0, len(flat), 5))


def blob_name(path, content_sha256, keys):
    mac = hmac.new(keys.name, path.encode("utf-8") + b"\x00" + content_sha256, hashlib.sha256).digest()
    return b32encode(mac[:20])


# MARK: the vault on disk

class BlobHeader:
    def __init__(self, plain):
        try:
            fields = json.loads(plain)
            self.version = fields["version"]
            self.path = fields["path"]
            self.size = fields["size"]
            self.mtime = fields["mtime"]
            self.sha256 = base64.b64decode(fields["contentSHA256"])
            self.content_key = base64.b64decode(fields["contentKey"])
        except (ValueError, KeyError, TypeError):
            raise FormatError("malformed blob header")
        if self.version != FORMAT_VERSION:
            raise FormatError("unsupported format version %r" % (self.version,))
        if len(self.content_key) != 32 or not isinstance(self.size, int) or self.size < 0:
            raise FormatError("malformed blob header")


class Vault:
    def __init__(self, path):
        if os.path.isdir(os.path.join(path, ROOT)):
            self.root = os.path.join(path, ROOT)
        elif os.path.basename(os.path.normpath(path)) == ROOT and os.path.isdir(path):
            self.root = path
        else:
            raise FormatError("no %s folder at %s" % (ROOT, path))

    # MARK: indexes

    def generations(self):
        found = []
        for name in os.listdir(self.root):
            match = MANIFEST_NAME.match(name)
            if match:
                found.append(int(match.group(1)))
        return sorted(found)

    def open_index(self, generation, keys):
        with open(os.path.join(self.root, "m-%010d" % generation), "rb") as handle:
            box = handle.read()
        plain = open_box(keys.manifest, box, struct.pack(">Q", generation))
        try:
            index = json.loads(plain)
        except ValueError:
            raise FormatError("malformed snapshot index")
        if index.get("version") != FORMAT_VERSION:
            raise FormatError("unsupported format version %r" % (index.get("version"),))
        if index.get("generation") != generation:
            raise FormatError("snapshot index names another generation")
        return index

    def indexes(self, keys):
        """every readable index oldest first, and the generations that failed"""
        readable, damaged = [], []
        for generation in self.generations():
            try:
                readable.append(self.open_index(generation, keys))
            except (FormatError, OSError):
                damaged.append(generation)
        return readable, damaged

    # MARK: blobs

    def blob_names(self):
        directory = os.path.join(self.root, "b")
        if not os.path.isdir(directory):
            return []
        return sorted(name for name in os.listdir(directory) if BLOB_NAME.match(name))

    def _read_header(self, handle, name, keys):
        length_word = handle.read(4)
        if len(length_word) != 4:
            raise FormatError("not a blob: too short")
        header_length = struct.unpack(">I", length_word)[0]
        if header_length < SEAL_OVERHEAD or header_length > MAX_HEADER_BOX:
            raise FormatError("not a blob: implausible header length")
        box = handle.read(header_length)
        if len(box) != header_length:
            raise FormatError("truncated blob header")
        return BlobHeader(open_box(keys.header, box, name.encode("utf-8") + b"h")), header_length

    def read_header(self, name, keys):
        with open(os.path.join(self.root, "b", name), "rb") as handle:
            return self._read_header(handle, name, keys)[0]

    def open_blob(self, name, keys):
        """the header, and a generator of verified plaintext segments"""
        handle = open(os.path.join(self.root, "b", name), "rb")
        try:
            header, header_length = self._read_header(handle, name, keys)
            # the authenticated size dictates the layout exactly, so
            # truncation and trailing bytes are told apart before any content
            segment_count = max(1, -(-header.size // SEGMENT_SIZE))
            final_plain = header.size - (segment_count - 1) * SEGMENT_SIZE
            expected = 4 + header_length + header.size + SEAL_OVERHEAD * segment_count
            actual = os.fstat(handle.fileno()).st_size
            if actual < expected:
                raise FormatError("truncated blob")
            if actual > expected:
                raise FormatError("not a blob: trailing bytes")
        except Exception:
            handle.close()
            raise

        def segments():
            with handle:
                for index in range(segment_count):
                    final = index == segment_count - 1
                    wire = (final_plain if final else SEGMENT_SIZE) + SEAL_OVERHEAD
                    box = handle.read(wire)
                    if len(box) != wire:
                        raise FormatError("truncated blob")
                    aad = name.encode("utf-8") + struct.pack(">Q", index) + (b"\x01" if final else b"\x00")
                    yield open_box(header.content_key, box, aad)

        return header, segments()

    def rescue_index(self, keys):
        """a best-guess index from blob headers alone: latest mtime wins per path"""
        by_path = {}
        unreadable = 0
        for name in self.blob_names():
            try:
                header = self.read_header(name, keys)
            except (FormatError, OSError):
                unreadable += 1
                continue
            held = by_path.get(header.path)
            if held is not None and held[1].mtime >= header.mtime:
                continue
            by_path[header.path] = (name, header)
        entries = [{"path": header.path, "objectName": name, "size": header.size, "mtime": header.mtime,
                    "contentSHA256": base64.b64encode(header.sha256).decode("ascii")}
                   for name, header in by_path.values()]
        entries.sort(key=lambda entry: entry["path"])
        now = datetime.datetime.now(datetime.timezone.utc).timestamp() - REFERENCE_DATE
        index = {"version": FORMAT_VERSION, "generation": 0, "created": now,
                 "entries": entries, "carriedForwardFolders": []}
        return index, unreadable


# MARK: restoring

def safe_components(path):
    parts = path.split("/")
    if len(parts) < 2 or any(part in ("", ".", "..") for part in parts):
        return None
    return parts


def check_blob(vault, keys, entry, deliver):
    """open the entry's blob, verify every segment and hand the plaintext to
    deliver; the failure reason as a string, or None"""
    try:
        header, segments = vault.open_blob(entry["objectName"], keys)
    except FileNotFoundError:
        return "blob missing"
    except (FormatError, OSError) as error:
        return str(error)
    expected = base64.b64decode(entry["contentSHA256"])
    if header.path != entry["path"] or header.sha256 != expected:
        return "the blob's header disagrees with the index"
    digest = hashlib.sha256()
    try:
        for segment in segments:
            digest.update(segment)
            deliver(segment)
    except (FormatError, OSError) as error:
        return str(error)
    if digest.digest() != expected:
        # every tag verified, so this is what was backed up; the file most
        # likely changed while it was being read. Kept, and said.
        print("  warning: content differs from the hash in the index: %s" % entry["path"])
    return None


def restore_entry(vault, keys, entry, output):
    parts = safe_components(entry["path"])
    if parts is None:
        return "unsafe path"
    destination = os.path.join(output, *parts)
    partial = destination + ".nas-reborn-partial"
    try:
        os.makedirs(os.path.dirname(destination), exist_ok=True)
        with open(partial, "wb") as out:
            reason = check_blob(vault, keys, entry, out.write)
        if reason is None:
            mtime = unix_time(entry["mtime"])
            os.utime(partial, (mtime, mtime))
            os.replace(partial, destination)
    except OSError as error:
        reason = "cannot write: %s" % error
    if reason is not None:
        try:
            os.remove(partial)
        except OSError:
            pass
    return reason


def select_entries(index, prefix):
    entries = index["entries"]
    if prefix:
        prefix = prefix.rstrip("/")
        entries = [entry for entry in entries
                   if entry["path"] == prefix or entry["path"].startswith(prefix + "/")]
    return entries


def report(verb, done, total_bytes, failures):
    print("%s %d files, %s" % (verb, done, human(total_bytes)))
    for path, reason in failures:
        print("  failed: %s (%s)" % (path, reason))
    return 1 if failures else 0


def restore_index(vault, keys, index, output, prefix=None):
    entries = select_entries(index, prefix)
    if not entries:
        sys.exit("nothing in this snapshot matches %r" % prefix)
    os.makedirs(output, exist_ok=True)
    done, total_bytes, failures = 0, 0, []
    for entry in entries:
        reason = restore_entry(vault, keys, entry, output)
        if reason is None:
            print("  " + entry["path"])
            done += 1
            total_bytes += entry["size"]
        else:
            failures.append((entry["path"], reason))
    return report("restored", done, total_bytes, failures)


# MARK: commands

def read_keys(args):
    if args.phrase_file:
        with open(args.phrase_file, "r", encoding="utf-8") as handle:
            phrase = handle.read()
    elif not sys.stdin.isatty():
        phrase = sys.stdin.readline()
    else:
        phrase = getpass.getpass("Recovery phrase: ")
    try:
        return Keys.from_phrase(phrase)
    except ValueError as error:
        sys.exit("error: %s" % error)


def open_vault(args):
    try:
        return Vault(args.vault)
    except FormatError as error:
        sys.exit("error: %s" % error)


def choose_index(vault, keys, generation):
    readable, damaged = vault.indexes(keys)
    if not readable:
        if damaged:
            sys.exit("error: no snapshot index opens with this phrase (%d present). "
                     "Check the phrase; if it is right the indexes are damaged and "
                     "\"rescue\" can rebuild one from the blobs." % len(damaged))
        sys.exit("error: this vault holds no snapshot index; \"rescue\" can rebuild one from the blobs")
    if generation is None:
        return readable[-1]
    for index in readable:
        if index["generation"] == generation:
            return index
    sys.exit("error: no readable snapshot %d; \"list\" shows the ones there are" % generation)


def command_list(args):
    vault = open_vault(args)
    keys = read_keys(args)
    readable, damaged = vault.indexes(keys)
    if not readable:
        choose_index(vault, keys, None)   # exits with the right message
    print("%8s  %-16s  %8s  %10s" % ("snapshot", "created", "files", "size"))
    for index in readable:
        note = ""
        if index.get("carriedForwardFolders"):
            note = "  folders carried forward: " + ", ".join(index["carriedForwardFolders"])
        print("%8d  %-16s  %8d  %10s%s" % (index["generation"], format_date(index["created"]),
                                           len(index["entries"]),
                                           human(sum(entry["size"] for entry in index["entries"])), note))
    if damaged:
        print("%d damaged snapshot index(es) skipped: %s" % (len(damaged), ", ".join(map(str, damaged))))
    return 0


def command_restore(args):
    vault = open_vault(args)
    keys = read_keys(args)
    index = choose_index(vault, keys, args.snapshot)
    print("restoring snapshot %d from %s into %s" % (index["generation"], format_date(index["created"]),
                                                     args.output))
    return restore_index(vault, keys, index, args.output, args.path)


def command_verify(args):
    vault = open_vault(args)
    keys = read_keys(args)
    index = choose_index(vault, keys, args.snapshot)
    print("verifying snapshot %d from %s" % (index["generation"], format_date(index["created"])))
    done, total_bytes, failures = 0, 0, []
    for entry in index["entries"]:
        reason = check_blob(vault, keys, entry, lambda segment: None)
        if reason is None:
            done += 1
            total_bytes += entry["size"]
        else:
            failures.append((entry["path"], reason))
    return report("verified", done, total_bytes, failures)


def command_rescue(args):
    vault = open_vault(args)
    keys = read_keys(args)
    index, unreadable = vault.rescue_index(keys)
    print("rebuilt an index of %d files from %d blobs; %d blobs did not open"
          % (len(index["entries"]), len(vault.blob_names()), unreadable))
    if not index["entries"]:
        sys.exit("error: no blob opens with this phrase")
    print("this is the latest version of every file ever backed up, not a snapshot")
    return restore_index(vault, keys, index, args.output)


def command_self_test(args):
    # RFC 4648 vectors, lowercase and unpadded
    for plain, encoded in ((b"", ""), (b"f", "my"), (b"fo", "mzxq"), (b"foo", "mzxw6"),
                           (b"foob", "mzxw6yq"), (b"fooba", "mzxw6ytb"), (b"foobar", "mzxw6ytboi")):
        assert b32encode(plain) == encoded, encoded
        assert b32decode(encoded) == plain, encoded
    # the frozen naming vector shared with the app's own tests: it pins the
    # HKDF info strings, the HMAC message layout, the 160-bit truncation and
    # the alphabet at once
    keys = Keys(bytes([7]) * 32)
    name = blob_name("photos/cat.jpg", hashlib.sha256(b"hello").digest(), keys)
    assert name == "ifdirknpiksudvpj3iw3j4d2qtuvyb6h", name
    # the phrase round-trips, forgives transcription noise and catches typos
    phrase = keys.phrase()
    assert len(phrase.replace("-", "")) == 55 and phrase.count("-") == 10, phrase
    assert Keys.from_phrase(phrase.upper().replace("-", " ") + "\n").master == keys.master
    broken = list(phrase)
    broken[3] = "a" if broken[3] != "a" else "b"
    try:
        Keys.from_phrase("".join(broken))
        raise AssertionError("a typo went unnoticed")
    except ValueError:
        pass
    # a box sealed here opens here, and not under a different AAD
    box = AESGCM(keys.header).encrypt(bytes(12), b"plain", b"aad")
    assert open_box(keys.header, bytes(12) + box, b"aad") == b"plain"
    try:
        open_box(keys.header, bytes(12) + box, b"other")
        raise AssertionError("a wrong AAD went unnoticed")
    except FormatError:
        pass
    print("self-test passed")
    return 0


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="restore a NAS Reborn backup without the app",
        epilog="the module docstring holds the vault format; read it with pydoc")
    commands = parser.add_subparsers(dest="command", metavar="COMMAND")
    commands.required = True

    def add(name, help_text, output=False, snapshot=False, path=False):
        command = commands.add_parser(name, help=help_text)
        command.add_argument("vault", help="folder holding nas-reborn-backup, or that folder")
        if output:
            command.add_argument("output", help="folder to write the files into")
        if snapshot:
            command.add_argument("--snapshot", type=int, metavar="N", help="generation, default the newest")
        if path:
            command.add_argument("--path", metavar="PREFIX", help="only this folder or file")
        command.add_argument("--phrase-file", metavar="FILE", help="read the recovery phrase from FILE")
        return command

    add("list", "show the snapshots in the vault").set_defaults(run=command_list)
    add("restore", "write a snapshot to a folder", output=True, snapshot=True,
        path=True).set_defaults(run=command_restore)
    add("verify", "decrypt a snapshot without writing", snapshot=True).set_defaults(run=command_verify)
    add("rescue", "rebuild from blob headers when every index is lost",
        output=True).set_defaults(run=command_rescue)
    commands.add_parser("self-test", help="check this script against the frozen vectors") \
        .set_defaults(run=command_self_test)

    args = parser.parse_args(argv)
    return args.run(args)


if __name__ == "__main__":
    sys.exit(main())