kc-apt-file

#!/usr/bin/python3

# Copyright © 2011-2024 Jakub Wilk <jwilk@jwilk.net>
# SPDX-License-Identifier: MIT

'''
apt-file replacement with Tokyo Cabinet backend
'''

0_0  # Python >= 3.6 is required

import argparse
import collections
import gzip
import os
import re
import shutil
import subprocess as ipc
import sys
import tempfile
import zlib

import apt
import apt_pkg

import kyotocabinet as kc

class xdg:
    cache_home = os.environ.get('XDG_CACHE_HOME') or ''
    if not os.path.isabs(cache_home):
        cache_home = os.path.join(os.path.expanduser('~'), '.cache')

def makedirs700(path):
    # TODO: Get rid of this function once
    # https://github.com/python/cpython/issues/86533
    # ("Restore os.makedirs ability to apply mode to all directories created")
    # is fixed.
    if os.path.isdir(path):
        return
    parent = os.path.dirname(path)
    if parent:
        makedirs700(parent)
    try:
        os.mkdir(path, 0o700)
    except OSError:
        if not os.path.isdir(path):
            raise

def get_cache_filename(*, mkdir=True):
    home = xdg.cache_home
    path = os.path.join(home, 'debian')
    if mkdir:
        makedirs700(path)
    return os.path.join(path, 'apt-file.kch')

class AcquireProgress(apt.progress.text.AcquireProgress):
    def __init__(self):
        apt.progress.text.AcquireProgress.__init__(self, outfile=sys.stderr)

class AcquireFailed(RuntimeError):
    pass

class Acquire(apt_pkg.Acquire):

    def run(self):
        rc = apt_pkg.Acquire.run(self)
        if rc != self.RESULT_CONTINUE:
            raise AcquireFailed('fetching files failed')
        for file in self.items:
            if file.status != file.STAT_DONE:
                raise AcquireFailed(f'fetching file failed: {file.desc_uri}')

class ParseError(RuntimeError):
    pass

def do_update(options):
    parse = re.compile(br'^(.*\S)\s+(\S+)$').match
    data = collections.defaultdict(list)
    if options.quiet:
        fetcher = Acquire()
    else:
        progress = AcquireProgress()
        fetcher = Acquire(progress)
    files = {}
    tmpdir = tempfile.mkdtemp(prefix='kc-apt-file.')
    try:
        mirror = options.mirror
        dist = options.distribution
        arch = options.architecture
        for section in options.sections:
            uri = f'{mirror}/dists/{dist}/{section}/Contents-{arch}.gz'
            files[section] = apt_pkg.AcquireFile(fetcher,
                uri=uri,
                descr=f'{dist}/{section}/Contents-{arch}',
                destfile=os.path.join(tmpdir, section),
            )
        fetcher.run()
        for file in files.values():
            uri = file.desc_uri
            skip = True
            with gzip.open(file.destfile) as file:
                for n, line in enumerate(file):
                    if n == 0 and not line.startswith(b'This file maps each file '):
                        skip = False
                    if skip:
                        if line.startswith(b'FILE'):
                            skip = False
                    else:
                        match = parse(line)
                        if match is None:
                            raise ParseError(f'cannot parse line: {repr(line)[1:]}')
                        filename, packages = match.groups()
                        for package in [pkg.split(b'/')[-1] for pkg in packages.split(b',')]:
                            data[package] += [filename]
            if skip:
                raise ParseError(f'{uri}: empty file')
    finally:
        shutil.rmtree(tmpdir)
    db = kc.DB()
    db.open(options.cache_file, db.OWRITER | db.OCREATE | db.OTRUNCATE)
    try:
        for package, filenames in data.items():
            db[package] = zlib.compress(b'\n'.join(filenames), 1)
            filenames[:] = []
    finally:
        db.close()

def strings_to_regexp(strings, *, regexp=False, whole=True):
    if regexp:
        result = '|'.join(strings)
    else:
        r_strings = str.join('|', map(re.escape, strings))
        result = f'(?:{r_strings})'
        if whole:
            result = f'^{result}$'
    return re.compile(result)

def do_list(options):
    regexp = strings_to_regexp(options.packages, regexp=options.regexp)
    matching = re.compile(regexp).search
    db = kc.DB()
    db.open(options.cache_file, db.OREADER)
    try:
        for package in db:
            package = package.decode('ASCII')
            if matching(package):
                filenames = zlib.decompress(db[package]).decode('UTF-8')
                filenames = filenames.splitlines()
                for filename in filenames:
                    print(f'{package}: /{filename}')
    finally:
        db.close()

def do_search(options):
    regexp = strings_to_regexp(options.terms, regexp=options.regexp, whole=False)
    matching = re.compile(regexp).search
    db = kc.DB()
    db.open(options.cache_file, db.OREADER)
    try:
        for package in db:
            filenames = zlib.decompress(db[package]).decode('UTF-8')
            filenames = filenames.splitlines()
            package = package.decode('ASCII')
            for filename in filenames:
                if matching('/' + filename):
                    print(f'{package}: /{filename}')
    finally:
        db.close()

class default_architecture:

    def __init__(self):
        self._cache = None

    def __str__(self):
        if self._cache is None:
            self._cache = ipc.getoutput('dpkg --print-architecture').strip()
        return self._cache

class default:
    mirror = 'https://deb.debian.org/debian'
    distribution = 'unstable'
    architecture = default_architecture()
    sections = ['main', 'contrib', 'non-free']

def unexpand_tilde(path):
    home = os.path.expanduser('~/')
    if path.startswith(home):
        path = '~/' + path[len(home):]
    return path

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    cache_filename = get_cache_filename(mkdir=False)
    parser.add_argument('--cache-file', metavar='FILE',
        help=f'use this cache file (default: {unexpand_tilde(cache_filename)})'
    )
    subparsers = parser.add_subparsers(dest='command')
    cmd_update = subparsers.add_parser('update', help='download indices from a Debian mirror')
    cmd_update.add_argument('-m', '--mirror', default=default.mirror,
        help=f'Debian mirror to use (default: {default.mirror})'
    )
    cmd_update.add_argument('-a', '--architecture', default=default.architecture,
        help='architecture to use'
    )
    cmd_update.add_argument('-d', '--distribution', default=default.distribution,
        help=f'distribution to use (default: {default.distribution})'
    )
    s_sections = str.join(', ', default.sections)
    cmd_update.add_argument('-s', '--sections', metavar='SECTION', nargs='+', default=default.sections,
        help=f'distribution to use (default: {s_sections})'
    )
    cmd_update.add_argument('-q', '--quiet', action='store_true')
    cmd_list = subparsers.add_parser('list', help='list files of the package(s)')
    cmd_list.add_argument('packages', metavar='PACKAGE', nargs='+')
    cmd_list.add_argument('-E', '--regexp', action='store_true')
    cmd_search = subparsers.add_parser('search', help='search for file(s)')
    cmd_search.add_argument('terms', metavar='TERM', nargs='+')
    cmd_search.add_argument('-E', '--regexp', action='store_true')
    options = parser.parse_args()
    if options.cache_file is None:
        options.cache_file = get_cache_filename()
    command = options.command
    globals()['do_' + command](options)

if __name__ == '__main__':
    main()

# vim:ts=4 sts=4 sw=4 et