cloud-vault

#!/usr/bin/env python3
import argparse, base64, contextlib, datetime, hashlib, json, os, \
       re, shlex, sqlite3, stat, subprocess, sys, urllib.parse
assert sys.hexversion >= 0x03050000 # Python 3.5 is required
sys.path.insert(0, os.path.expanduser("~/stuff/_urandom"))
import utils

EMPTY_HASH = "sha512:z4PhNX7vuL3xVChQ1m2AB9Yg5AUL"
DIGEST_LEN = 21 # in bytes
MIN_HASH_PREFIX = 8 # in bytes

class UserError(Exception):
    pass

def bihash_file(hasher1, hasher2, file):
    if hasher1 == hasher2:
        h = utils.hash_file(hasher1, file)
        return h, h
    h1 = hasher1()
    h2 = hasher2()
    for block in iter(lambda: file.read(1 << 20), b""):
        h1.update(block)
        h2.update(block)
    return h1, h2

def get_file_hash(file, known=None):
    hasher = hashlib.sha512
    if known is not None:
        known_hasher, known_hash = known
        assert len(known_hash) >= MIN_HASH_PREFIX * 2
    else:
        known_hasher = hasher
    with open(file, "rb") as f:
        hashed, known_hashed = bihash_file(hashlib.sha512, known_hasher, f)
    if known is not None:
        if not known_hashed.hexdigest().startswith(known_hash):
            raise Exception("hash mismatch:\n"
                            "found:    {!r}\n"
                            "provided: {!r}"
                            .format(known_hashed.hexdigest(),
                                    known_hash))
    return "sha512:" + display_digest(hashed)

def get_hasher(hash_algorithm):
    if hash_algorithm not in hashlib.algorithms_available:
        return None
    return getattr(hashlib, hash_algorithm)

def convert_hash_to_hex(hash):
    algorithm, digest = hash.split(":", 1)
    hex_digest = base64.b16encode(base64.urlsafe_b64decode(digest))
    return algorithm + ":" + hex_digest.decode("ascii").lower()

def display_digest(hashed):
    digest = hashed.digest()[:DIGEST_LEN]
    return base64.urlsafe_b64encode(digest).decode("ascii")

def display_mode(mode):
    return "755" if mode & stat.S_IXUSR else "644"

def parse_mode(mode):
    if mode not in ["644","755"]:
        raise ValueError("unsupported mode: {}".format(mode))
    return int(mode, base=8) & 0o777

def display_time(timestamp_ns):
    return (
        re.sub("[:-]", "", datetime.datetime.utcfromtimestamp(
            timestamp_ns // 1000000000
        ).isoformat()) +
        ".{:09}Z".format(timestamp_ns % 1000000000))

def parse_time(isotime):
    t, ns = isotime.split(".")
    t = (datetime.datetime.strptime(t, "%Y%m%dT%H%M%S") -
         datetime.datetime(1970, 1, 1)) // datetime.timedelta(seconds=1)
    ns = int(ns.rstrip("Z"))
    return t * 1000000000 + ns

def generate_remote_id(seed, hash):
    if seed is None:
        return None
    return display_digest(hashlib.sha512((seed + hash).encode("utf-8")))

def warn(s):
    sys.stderr.write(s)
    sys.stderr.flush()

def db_hash_exists(db, hash):
    return bool(tuple(db.execute("SELECT 1 FROM hashes WHERE hash = ?",
                                 [hash])))

def db_insert_hash(db, hash, size):
    with db:
        db.execute("BEGIN")
        db.execute("INSERT OR REPLACE INTO hashes (hash, size) VALUES (?, ?)",
                   [hash, size])

def db_remove_hash(db, hash):
    with db:
        db.execute("BEGIN")
        db.execute("DELETE FROM hashes WHERE hash = ?", [hash])

def db_all_hashes(db):
    return tuple(x for x, in db.execute("SELECT hash FROM hashes"))

def db_get_size_by_hash(db, hash):
    for size, in db.execute("SELECT size FROM hashes WHERE hash = ?", [hash]):
        return size
    return None

def transfer_rm(transfer, hash, remote):
    # handle empty files specially because cloud-transfer doesn't remember
    # empty files at all
    if hash == EMPTY_HASH:
        pass
    else:
        subprocess.check_call(transfer + ["rm", "--", remote])

def transfer_download(transfer, hash, remote, local):
    # handle empty files specially because cloud-transfer doesn't remember
    # empty files at all
    if hash == EMPTY_HASH:
        with open(local, "w"):
            pass
    else:
        try:
            subprocess.check_call(transfer + ["download", "--", remote, local])
        except:
            try:
                if os.lstat(local).st_size == 0:
                    os.remove(local)
            except OSError:
                pass
            raise

def transfer_upload(transfer, hash, local, remote):
    hex_hash = convert_hash_to_hex(hash)
    subprocess.check_call(transfer + ["upload", "--hash", hex_hash, "--",
                                      local, remote])

def is_vault_link(file):
    return (os.path.islink(file) and
            os.readlink(file).startswith("vault://"))

def get_info_from_vault_link(db, seed, file):
    target = os.readlink(file)
    hash, qs = re.match(r"vault://([^?]+)\?(.*)", target).groups()
    attrs = dict(urllib.parse.parse_qsl(qs))
    return {
        "hash": hash,
        "hex_hash": convert_hash_to_hex(hash),
        "remote_id": generate_remote_id(seed, hash),
        "mode": attrs["mode"],
        "mtime": attrs["mtime"],
        "size": db_get_size_by_hash(db, hash) if db is not None else None,
    }

def get_info(seed, file, ignore_mode, known=None):
    st = os.lstat(file)
    if not stat.S_ISREG(st.st_mode):
        raise UserError("not a file: {!r}".format(file))
    hash = get_file_hash(file, known=known)
    return {
        "hash": hash,
        "hex_hash": convert_hash_to_hex(hash),
        "remote_id": generate_remote_id(seed, hash),
        "mode": display_mode(0o644 if ignore_mode else st.st_mode),
        "mtime": display_time(st.st_mtime_ns),
        "size": st.st_size,
    }

def info(files, ignore_mode, seed, transfer, db, **kwargs):
    for file in files:
        if not os.path.lexists(file):
            raise UserError("does not exist: {}".format(file))
        if os.path.islink(file):
            info = get_info_from_vault_link(db, seed, file)
        else:
            info = get_info(seed, file, ignore_mode)
        sys.stdout.write(json.dumps(info, **utils.JSON_PRETTY) + "\n")

def get(files, new, seed, transfer, db, **kwargs):
    for file in files:
        if not os.path.lexists(file):
            raise UserError("does not exist: {}".format(file))
        if os.path.islink(file):
            if not is_vault_link(file):
                raise UserError("cannot get a non-vault link: {}".format(file))
        else:
            warn("skipping (already got): {}\n".format(file))
            continue
        info = get_info_from_vault_link(db, seed, file)
        mode = parse_mode(info["mode"])
        mtime = parse_time(info["mtime"])
        if not new and not db_hash_exists(db, info["hash"]):
            raise UserError("hash not in index: {} (use --new to "
                            "record a new hash)".format(info["hash"]))
        tmp = file + ".download.tmp" # keep name consistent to allow resumption
        transfer_download(transfer, info["hash"], info["remote_id"], tmp)
        tmp_hash = get_file_hash(tmp)
        if tmp_hash != info["hash"]:
            raise UserError("hash of downloaded file does not match:\n"
                            "downloaded: {}\n"
                            "link:       {}"
                            .format(tmp_hash, info["hash"]))
        db_insert_hash(db, info["hash"], os.path.getsize(tmp))
        os.utime(tmp, ns=(os.lstat(tmp).st_atime_ns, mtime))
        os.chmod(tmp, mode)
        os.rename(file, file + ".link")
        os.rename(tmp, file)

def put(files, delete, check, force, ignore_mode, seed, transfer, db, **kwargs):
    if check is not None:
        # determine checksum algorithm
        m = re.search("(\w*)sums?$", check.lower())
        if not m:
            raise UserError("could not determine checksum algorithm from "
                            "file extension")
        ck_algorithm, = m.groups()
        ck_hasher = get_hasher(ck_algorithm)
        if ck_hasher is None:
            raise UserError("unsupported checksum algorithm: {}"
                            .format(ck_algorithm))
        # populate checksums dictionary
        checksums = {}
        with open(check) as f:
            for line in f:
                if not line.strip():
                    continue
                m = re.match("([0-9a-fA-F]+)  (.*)", line)
                if not m:
                    raise UserError("could not parse line in checksum file: {}"
                                    .format(line))
                ck_hash, ck_path = m.groups()
                ck_path = os.path.realpath(
                    os.path.join(os.path.dirname(check), ck_path))
                if ck_path in checksums:
                    raise UserError("checksum file contains conflicting hashes "
                                    "for: {}".format(ck_path))
                checksums[ck_path] = ck_hash

        # bail early if one of the checksums is missing
        for file in files:
            if os.path.realpath(file) not in checksums:
                raise UserError("no entry in checksum file for: {}"
                                .format(file))

    for file in files:
        if is_vault_link(file):
            warn("skipping (already put): {}\n".format(file))
            continue
        if os.path.islink(file):
            raise UserError("cannot put a non-vault link: {}".format(file))
        bak_file = file + ".bak"
        if not force and os.path.exists(bak_file):
            raise UserError("backup file already exists: {}".format(bak_file))
        if check is None:
            known = None
        else:
            known = ck_hasher, checksums[os.path.realpath(file)]
        info = get_info(seed, file, ignore_mode, known=known)
        if db_hash_exists(db, info["hash"]):
            warn("skipping upload (already indexed): {}\n"
                 .format(file))
        else:
            remote_id = generate_remote_id(seed, info["hash"])
            transfer_upload(transfer, info["hash"], file, info["remote_id"])
        db_insert_hash(db, info["hash"], info["size"])
        target = "vault://{}?{}".format(info["hash"], urllib.parse.urlencode([
            ("mode", info["mode"]),
            ("mtime", info["mtime"]),
        ]))
        if not force:
            with open(bak_file, "x"):
                pass
        os.rename(file, bak_file)
        os.symlink(target, file)
        if delete:
            os.remove(bak_file)

def gc_mark(hashes, file):
    if not is_vault_link(file):
        return
    hash = get_info_from_vault_link(None, None, file)["hash"]
    try:
        hashes.remove(hash)
    except KeyError:
        pass

def gc(paths, purge, seed, transfer, db, purge_index, **kwargs):
    hashes = set(db_all_hashes(db))
    for path in paths:
        if os.path.isdir(path):
            for (dirpath, _, filenames) in os.walk(path):
                for filename in filenames:
                    gc_mark(hashes, os.path.join(dirpath, filename))
        else:
            gc_mark(hashes, path)
    for hash in sorted(hashes):
        remote_id = generate_remote_id(seed, hash)
        sys.stdout.write('{{"hash": "{}", "remote_id": "{}"}}\n'
                         .format(hash, remote_id))
        if purge:
            try:
                transfer_rm(transfer, hash, remote_id)
                db_remove_hash(db, hash)
            except subprocess.CalledProcessError as e:
                warn(str(e) + "\n")
        elif purge_index:
            db_remove_hash(db, hash)

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--index", required=True,
                   help="SQLite3 database to store the known hashes")
    p.add_argument("--seed", required=True,
                   help="seed for generating remote IDs")
    p.add_argument("--transfer", required=True, type=shlex.split,
                   help="cloud-transfer command")
    p.add_argument("--ignore-mode", action="store_true",
                   help="Treat the file mode as 644")
    sp = p.add_subparsers()

    spp = sp.add_parser("gc")
    spp.set_defaults(func=gc)
    spp.add_argument("paths", nargs="+")
    spp.add_argument("--purge", action="store_true",
                     help="Delete the files at remote (implies --purge-index)")
    spp.add_argument("--purge-index", action="store_true",
                     help="Delete the files in index only")

    spp = sp.add_parser("info", help="Retrieve information about file/link")
    spp.set_defaults(func=info)
    spp.add_argument("files", nargs="+", type=os.path.abspath,
                     help="filename of the local file")

    spp = sp.add_parser("get",
                        help=("Replace a vault link with the actual file. "
                              "The link is backed up to *.link"))
    spp.set_defaults(func=get)
    spp.add_argument("--new", action="store_true",
                     help=("Add a new file that has already been uploaded "
                           "but is not in index.  The file will be stored "
                           "in index if download and verification succeed."))
    spp.add_argument("files", nargs="+", type=os.path.abspath,
                     help="path to the local metadata files")

    spp = sp.add_parser("put",
                        help=("Upload and replace the file with a vault link. "
                              "The original file is backed up to *.bak"))
    spp.set_defaults(func=put)
    spp.add_argument("--delete", action="store_true",
                     help="Don't make a backup of the file.")
    spp.add_argument("-c", "--check",
                     help="Verify files using the provided checksum file")
    spp.add_argument("-f", "--force", action="store_true",
                     help="Overwrite the backup file if it already exists.")
    spp.add_argument("files", nargs="+", type=os.path.abspath,
                     help="filename of the local file")

    kwargs = vars(p.parse_args())
    func = kwargs.pop("func", None)
    if not func:
        p.print_usage()
        return 2

    os.makedirs(os.path.realpath(os.path.dirname(kwargs["index"])),
                exist_ok=True)
    with contextlib.closing(sqlite3.connect(kwargs["index"])) as db:
        migrations = [
            [
                """
                CREATE TABLE hashes
                ( hash TEXT UNIQUE
                )
                """,
                """
                CREATE TABLE aliases
                ( alias TEXT UNIQUE
                , hash TEXT
                , FOREIGN KEY(hash) REFERENCES hashes(hash)
                )
                """,
            ],
            [
                """
                ALTER TABLE "hashes" ADD COLUMN "size" INTEGER;
                """,
            ],
        ]
        db.execute("PRAGMA auto_vacuum = FULL")
        with db:
            db.execute("BEGIN")
            version = max(0, tuple(db.execute("PRAGMA user_version"))[0][0])
            if version > len(migrations):
                raise Exception("database is newer than current version")
            for migration in migrations[version:]:
                for stmt in migration:
                    db.execute(stmt)
            db.execute("PRAGMA user_version = {}".format(len(migrations)))
        try:
            return func(db=db, **kwargs) or 0
        except UserError as e:
            sys.stderr.write("cloud-vault: {}\n".format(e))
            sys.stderr.flush()
            sys.exit(1)

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