patch_all_nix_elfs.py

"""
Replace all NEEDEDs of all ELF files in /nix/store with absolute paths.
Requires python-magic.
"""

from __future__ import print_function
import os.path
import subprocess
from collections import namedtuple
import functools

import magic


class FileType(namedtuple('FileType', ['base_type', 'link_type'])):
    NOT_INTERESTING = 'not_interesting'
    EXECUTABLE = 'executable'
    SHARED_OBJECT = 'shared_object'
    SYMLINK = 'symlink'

    NA = 'na'
    STATIC = 'static'
    DYNAMIC = 'dynamic'

    @classmethod
    def from_path(cls, filepath):
        file_output = magic.from_file(filepath)

        if file_output.startswith('ELF '):
            link_type = cls.DYNAMIC
            if 'statically linked' in file_output:
                link_type = cls.STATIC

            if 'executable' in file_output:
                return cls(cls.EXECUTABLE, link_type)
            if 'shared object' in file_output:
                return cls(cls.SHARED_OBJECT, link_type)

            return cls(cls.NOT_INTERESTING, cls.NA)

        if file_output.startswith('symbolic link '):
            return cls(cls.SYMLINK, cls.NA)

        return cls(cls.NOT_INTERESTING, cls.NA)


class StoreFile(namedtuple('StoreFile', ['path', 'file_type', 'neededs', 'rpaths'])):
    @classmethod
    def from_path(cls, path):
        file_type = FileType.from_path(path)
        if file_type.base_type == FileType.NOT_INTERESTING:
            return cls(path, file_type, [], [])

        neededs = []
        if file_type.link_type == FileType.DYNAMIC:
            neededs = subprocess.check_output(['patchelf', '--print-needed', path]).strip().splitlines()

        rpaths = []
        if file_type.link_type == FileType.DYNAMIC:
            rpaths = subprocess.check_output(['patchelf', '--print-rpath', path]).strip().split(':')

        return cls(path, file_type, neededs, rpaths)


    def has_needed(self, needed, check_basename=False):
        for line in self.neededs:
            if check_basename:
                if os.path.basename(line) == needed:
                    return True
            elif line == needed:
                return True
        return False

    def iter_basename_neededs(self):
        for needed in self.neededs:
            if os.path.basename(needed) == needed:
                yield needed


# TODO: update target_file.neededs as well?
def replace_needed(target_file, old_needed, new_needed):
    if target_file.has_needed(old_needed):
        subprocess.check_call(['patchelf', '--replace-needed', old_needed, new_needed, target_file.path])
        print("  {old_needed} -> {new_needed}".format(**locals()))
    else:
        print("  {old_needed} not needed; skipping".format(**locals()))


def add_needed(target_file, needed):
    if target_file.has_needed(needed, check_basename=True):
        print("  {needed} already in needed; skipping".format(**locals()))
    else:
        subprocess.check_call(['patchelf', '--add-needed', needed, target_file.path])
        print("  {needed} added".format(**locals()))


NOT_INTERESTING_EXTENSIONS = [
    '.drv',
    '.gz',
    '.h',
    '.hh',
    '.hpp',
    '.hs',
    '.java',
    '.mo',
    '.nix',
    '.patch',
    '.pod',
    '.py',
    '.pl',
    '.rb',
    '.sh',
    '.txt',
    '.yaml',
    '.yml',
    '.xml',
]


def heuristically_not_interesting(path):
    _, ext = os.path.splitext(path)
    if ext in NOT_INTERESTING_EXTENSIONS:
        return True
    return False


def iter_dynamic_elf_files(store_root):
    paths = subprocess.check_output(['find', store_root, '-xtype', 'f']).splitlines()
    for filepath in paths:
        if heuristically_not_interesting(filepath):
            continue
        print(filepath)
        store_file = StoreFile.from_path(filepath)
        if store_file.file_type.base_type == FileType.NOT_INTERESTING:
            continue
        if store_file.file_type.link_type != FileType.DYNAMIC:
            continue
        yield store_file


def chmod_r_in(parent_dir, mode):
    print('chmod -R {mode} {parent_dir}/*'.format(**locals()))
    for member in os.listdir(parent_dir):
        subprocess.check_call(['chmod', '-R', mode, os.path.join(parent_dir, member)])


def with_writable_store(fn):
    @functools.wraps(fn)
    def wrapper(store_root):
        try:
            chmod_r_in(store_root, 'u+w')
            fn(store_root)
        finally:
            chmod_r_in(store_root, 'u-w')
    return wrapper


class FileChecker(object):
    def __init__(self):
        # path -> bool
        # if a path is in this dict, it has been checked before.
        # if the value is true, the file existed when it was checked.
        # if the value is false, it did not.
        # if a path is not in this dict, it has not been checked before.
        self._checked_files = dict()

    def exists(self, path):
        if path not in self._checked_files:
            self._checked_files[path] = os.path.isfile(path)
        return self._checked_files[path]
    

@with_writable_store
def patch_all(store_root):
    checker = FileChecker()

    print('reading store files in', store_root)
    for elf in iter_dynamic_elf_files(store_root):
        for needed in elf.iter_basename_neededs():
            for rpath in elf.rpaths:
                abspath = os.path.join(rpath, needed)
                if checker.exists(abspath):
                    replace_needed(elf, needed, abspath)
                    break


if __name__ == '__main__':
    patch_all('/nix/store')