apt-find-foreign

#!/usr/bin/env python
# apt-find-foreign (part of ossobv/vcutil) // wdoekes/2018,2020,2022-2023
#   // Public Domain
"""
Script to find Debian/Ubuntu packages from "foreign" archives.

Outputs a list of apt source lists where the currently installed
packages come from. For lists that have fewer than 100 items, it also
lists the package names.

Usage example:

    apt-find-foreign

Output example:

    Lists with corresponding package counts:
      269     (local only)
      1       (updates available)
      3075    http://ubuntuserver.com/ubuntu
      2       http://custom-ppa.com/ubuntu

    Lists with very few packages (or with remarks):
      (local only)
        - libswresample-ffmpeg1
        - libquvi-scripts
        - qml-module-ubuntu-performancemetrics
        ...
      (updates available)
        - gpg
      http://custom-ppa.com/ubuntu
        - packagex
        - packagey

Another usage example:

    apt-get remove --purge `apt-find-foreign --rc --local | grep -v SOMETHING`

Because apt-cache policy has to do a lot of work, it takes more than 2 seconds
to run on my desktop.
"""
from argparse import ArgumentParser
from collections import OrderedDict, defaultdict
from subprocess import check_output
from tempfile import TemporaryFile

# Show entire package list for repo if there are at most this many:
VERY_FEW_PACKAGES_LTE = 40

CAT_LOCAL_ONLY = '(local only)'
CAT_RC_ONLY = '(rc only)'
CAT_DOWNGRADES_AVAILABLE = '(downgrades available)'
CAT_UPDATES_AVAILABLE = '(updates available)'
CAT_NON_URL = (
    CAT_LOCAL_ONLY, CAT_RC_ONLY,
    CAT_DOWNGRADES_AVAILABLE, CAT_UPDATES_AVAILABLE)


def check_english_output(command, **kwargs):
    assert 'env' not in kwargs, kwargs
    # Wipe all of LANG, LC_ALL, LC_MESSAGES to get C locale.
    return check_output(command, env={}, **kwargs).decode('utf-8')


class VersionState(object):
    def __init__(self, version, list_urls, version_in_list=None,
                 version_direction=0):
        self.version = version
        self.list_urls = tuple(sorted(list_urls))
        self.version_direction = version_direction
        self.version_in_list = version_in_list


class Package(object):
    def __init__(self, name, state, policy=None):
        self.name = name
        self.state = state
        # By injecting the policy, we save huge amounts of time. Doing one
        # fork()/exec() per package would be horrendously slow.
        self.policy = policy

    def get_policy(self):
        """
        Return apt-cache policy PKG_NAME output.

        For example::

            vcutil:
              Installed: 1.11
              Candidate: 1.11
              Version table:
             *** 1.11 500
                  500 http://example.com/ubuntu bionic/osso amd64 Packages
                  500 http://example.com/ubuntu bionic/osso i386 Packages
                  100 /var/lib/dpkg/status
               1.10 500
                  500 http://example.com/ubuntu bionic/osso amd64 Packages
                  500 http://example.com/ubuntu bionic/osso i386 Packages
        """
        if self.policy is None:
            self.policy = check_english_output(
                ['apt-cache', 'policy', self.name])
        return self.policy

    def get_version_state(self):
        """
        Return dpkg source version state from the policy.

        For example::

            VersionState('1.2.3', ('http://example.com/ubuntu',), '1.2.4')
        """
        if not hasattr(self, '_version_state'):
            self._version_state = self._extract_source_urls(
                self.name, self.state, self.get_policy())
        return self._version_state

    @staticmethod
    def _extract_source_urls(name, state, policy):
        out = policy.split('\n')
        outiter = iter(out)
        for line in outiter:
            if line.startswith('  Version table:'):
                break

        local_version = None
        local_urls = set()
        list_version = None
        local_list_version_cmp = None

        current_version_prio = next(outiter).split()
        while current_version_prio:
            urls = set()

            # ' *** 3.20220809.0ubuntu0.22.04.1 500 (phased 10%)'
            if (len(current_version_prio) > 3 and
                    current_version_prio[-2] == '(phased' and
                    current_version_prio[-1].endswith(')')):
                current_version_prio = current_version_prio[0:-2]

            # ' *** 1.11 500'
            if len(current_version_prio) == 3:
                current, version, prio = current_version_prio
                assert current == '***', (name, current_version_prio)
                local_version = version
                current = True
            elif len(current_version_prio) == 2:
                version, prio = current_version_prio
                current = False
            else:
                assert False, (name, current_version_prio)

            # '      500 http://example.com/ubuntu bionic/osso amd64 Packages'
            # '      100 /var/lib/dpkg/status'
            for line in outiter:
                if not line.startswith(' ' * 7):  # 7 for prio >999 else 8
                    break
                url = line.split()[1]
                if url != '/var/lib/dpkg/status':
                    urls.add(url)
            else:
                line = ''  # at EOF, but complete the while
            if current:  # found the URLs with '***'
                local_urls = urls

            # We did not have a list_version yet, but here it is.
            if list_version is None and urls:
                list_version = version
                if local_version == list_version:
                    local_list_version_cmp = 0
                elif local_version:  # local is newer?
                    local_list_version_cmp = -1
                else:
                    local_list_version_cmp = 1

            # Short circuit in case we have enough info.
            if None not in (local_version, list_version):
                break

            # Restart loop with the next (older) version.
            current_version_prio = line.split()

        assert bool(local_version) ^ bool(state == 'rc'), (name, state, policy)
        # list_version and local_list_version_tmp are None when local_only
        return VersionState(
            local_version, list_urls=local_urls, version_in_list=list_version,
            version_direction=local_list_version_cmp)

    def __repr__(self):
        return self.name


def dpkg_l():
    """
    Return list of [(package_name, state), ...].
    """
    # Check our/default arch
    arch = check_english_output(['dpkg', '--print-architecture']).strip()
    colon_arch = ':{}'.format(arch)  # ":amd64"
    colon_arch_neg_len = -len(colon_arch)

    def without_default_arch(pkg):
        return pkg[0:colon_arch_neg_len] if pkg.endswith(colon_arch) else pkg

    # Get listing, condense to [(package, state)] list
    out = check_english_output(['dpkg', '-l'])
    out = out.split('\n')[1:]
    out = [i for i in out if i and not i.startswith(tuple('|+'))]
    out = [list(reversed(i.split()))[-2:] for i in out]       # ("name", "ii")
    out = [(without_default_arch(i[0]), i[1]) for i in out]   # drop ":amd64"
    return out


def apt_cache_policy(package_names):
    """
    Return apt_cache_policy dict, one string per package.
    """
    with TemporaryFile(mode='w+b') as fp:
        fp.write('\0'.join(package_names).encode('utf-8'))
        fp.seek(0)
        out = check_english_output(
            ['xargs', '-0', 'apt-cache', 'policy'], stdin=fp).strip()
    out = out.split('\n')
    buf = []
    for line in out:
        if not line.startswith(' '):
            if buf:
                blob = '\n'.join(buf)
                name = blob.split(':\n', 1)[0]
                yield (name, blob)
                buf = []
        buf.append(line)
    if buf:
        blob = '\n'.join(buf)
        name = blob.split(':\n', 1)[0]
        yield (name, blob)


def get_packages_per_list():
    packages = OrderedDict(dpkg_l())
    policies = dict(apt_cache_policy(packages.keys()))
    lists = defaultdict(list)

    for pkg_name in packages.keys():
        p = Package(
            pkg_name, state=packages[pkg_name], policy=policies[pkg_name])
        vstate = p.get_version_state()

        if p.state not in ('ii', 'hi'):
            assert p.state == 'rc', (p.name, p.state)
            lists[CAT_RC_ONLY].append(p)
        elif vstate.version_in_list is None:  # exists in NO list
            lists[CAT_LOCAL_ONLY].append(p)
        elif vstate.list_urls:  # exists in a list, may not be newest though
            for list_ in vstate.list_urls:
                lists[list_].append(p)
        elif vstate.version_direction == 1:  # local version is older
            lists[CAT_UPDATES_AVAILABLE].append(p)
        elif vstate.version_direction == -1:  # local version is newer(!)
            lists[CAT_DOWNGRADES_AVAILABLE].append(p)
        else:
            assert False, (p.name, vstate.version, vstate.version_in_list)

    return lists


def show():
    lists = get_packages_per_list()

    print('Lists with corresponding package counts:')
    show_details = False
    for key in sorted(lists.keys()):
        print('  {}\t{}'.format(len(lists[key]), key))
        show_details = (
            show_details or
            len(lists[key]) <= VERY_FEW_PACKAGES_LTE or
            key in CAT_NON_URL)

    if show_details:
        print('')
        print('Lists with very few packages (or with remarks):')
        for key in sorted(lists.keys()):
            if len(lists[key]) <= VERY_FEW_PACKAGES_LTE or key in CAT_NON_URL:
                print('  {}'.format(key))
                for package in lists[key]:
                    print('    - {}'.format(package.name))


def main():
    parser = ArgumentParser()
    parser.add_argument('--local-only', action='store_true')
    parser.add_argument('--rc-only', action='store_true')
    parser.add_argument('--downgrades-available', action='store_true')
    parser.add_argument('--updates-available', action='store_true')
    args = parser.parse_args()

    if args.local_only or args.rc_only:
        assert not args.downgrades_available, 'mutually exclusive'
        assert not args.updates_available, 'mutually exclusive'
        # apt-find-foreign --rc --local | xargs apt-get remove --purge
        lists = get_packages_per_list()
        if args.local_only:
            for package in lists[CAT_LOCAL_ONLY]:
                print(package.name)
        if args.rc_only:
            for package in lists[CAT_RC_ONLY]:
                print(package.name)
    elif args.downgrades_available:
        assert not any([
            args.local_only, args.rc_only, args.updates_available])
        lists = get_packages_per_list()
        for package in lists[CAT_DOWNGRADES_AVAILABLE]:
            vstate = package.get_version_state()
            print('{:22s}  {:22s} >> {:s}'.format(
                package.name, vstate.version, vstate.version_in_list))
    elif args.updates_available:
        assert not any([
            args.local_only, args.rc_only, args.downgrades_available])
        lists = get_packages_per_list()
        for package in lists[CAT_UPDATES_AVAILABLE]:
            vstate = package.get_version_state()
            print('{:22s}  {:22s} << {:s}'.format(
                package.name, vstate.version, vstate.version_in_list))
    else:
        assert not any([
            args.local_only, args.rc_only,
            args.downgrades_available, args.updates_available])
        show()


if __name__ == '__main__':
    main()