aif-ng/aif/system/locales.py

import configparser
import copy
import io
import logging
import os
import re
import shutil
import subprocess


_logger = logging.getLogger(__name__)


# TODO: time
_locale_re = re.compile(r'^(?!#\s|)$')
_locale_def_re = re.compile(r'([^.]*)[^@]*(.*)')


class Locale(object):
    def __init__(self, chroot_base, locales_xml):
        self.xml = locales_xml
        self.chroot_base = chroot_base
        self.syslocales = {}
        self.userlocales = []
        self.rawlocales = None
        self._localevars = configparser.ConfigParser()
        self._localevars.optionxform = str
        self._localevars['BASE'] = {}
        self._initVars()

    def _initVars(self):
        for l in self.xml.findall('locale'):
            locale = l.text.strip()
            self._localevars['BASE'][l.attrib['name'].strip()] = locale
            if locale not in self.userlocales:
                self.userlocales.append(locale)
        if not self.userlocales:
            self.userlocales = ['en_US', 'en_US.UTF-8']
        _logger.debug('Rendered locales: {0}'.format(dict(self._localevars['BASE'])))
        _logger.debug('Rendered user locales: {0}'.format(','.join(self.userlocales)))
        return(None)

    def _verify(self):
        localegen = os.path.join(self.chroot_base, 'etc', 'locale.gen')  # This *should* be brand new.
        with open(localegen, 'r') as fh:
            self.rawlocales = fh.read().splitlines()
        for idx, line in enumerate(self.rawlocales[:]):
            if _locale_re.search(line) or line.strip() == '':
                continue
            locale, charset = line.split()
            locale = locale.replace('#', '')
            self.syslocales[locale] = charset
            if locale in self.userlocales:
                # "Uncomment" the locale (self.writeConf() actually writes the change)
                self.rawlocales[idx] = '{0} {1}'.format(locale, charset)
        _logger.debug('Rendered system locales: {0}'.format(self.syslocales))
        userl = set(self.userlocales)
        sysl = set(self.syslocales.keys())
        missing_locales = (userl - sysl)
        if (userl - sysl):
            _logger.error('Specified locale(s) {0} that does not exist on the target system.'.format(missing_locales))
            raise ValueError('Missing locale(s)')
        return(None)

    def writeConf(self):
        # We basically recreate locale-gen in python here, more or less.
        self._verify()
        localegen = os.path.join(self.chroot_base, 'etc', 'locale.gen')
        localedbdir = os.path.join(self.chroot_base, 'usr', 'lib', 'locale')
        localesrcdir = os.path.join(self.chroot_base, 'usr', 'share', 'i18n')
        with open(localegen, 'w') as fh:
            fh.write('# Generated by AIF-NG.\n\n')
            fh.write('\n'.join(self.rawlocales))
            fh.write('\n')
        _logger.info('Wrote: {0}'.format(localegen))
        # If only the locale DB wasn't in a hopelessly binary format.
        # These destinations are built by the below subprocess call.
        for root, dirs, files in os.walk(localedbdir):
            for f in files:
                fpath = os.path.join(root, f)
                os.remove(fpath)
            for d in dirs:
                dpath = os.path.join(root, d)
                shutil.rmtree(dpath)
        _logger.debug('Pruned locale destination.')
        for locale in self.userlocales:
            lpath = os.path.join(localesrcdir, 'locales', locale)
            charset = self.syslocales[locale]
            if os.path.isfile(lpath):
                ldef_name = locale
            else:
                ldef_name = _locale_def_re.sub(r'\g<1>\g<2>', locale)
                lpath = os.path.join(localesrcdir, 'locales', ldef_name)
            env = copy.deepcopy(dict(os.environ))
            env['I18NPATH'] = localesrcdir
            _logger.debug('Invocation environment: {0}'.format(env))
            cmd = subprocess.run(['localedef',
                                  '--force',
                                  # These are overridden by a prefix env var.
                                  # '--inputfile={0}'.format(lpath),
                                  # '--charmap={0}'.format(os.path.join(localesrcdir, 'charmaps', charset)),
                                  '--inputfile={0}'.format(ldef_name),
                                  '--charmap={0}'.format(charset),
                                  '--alias-file={0}'.format(os.path.join(self.chroot_base,
                                                                         'usr', 'share', 'locale', 'locale.alias')),
                                  '--prefix={0}'.format(self.chroot_base),
                                  locale],
                                 stdout = subprocess.PIPE,
                                 stderr = subprocess.PIPE,
                                 env = env)
            _logger.info('Executed: {0}'.format(' '.join(cmd.args)))
            if cmd.returncode != 0:
                _logger.warning('Command returned non-zero status')
                _logger.debug('Exit status: {0}'.format(str(cmd.returncode)))
                for a in ('stdout', 'stderr'):
                    x = getattr(cmd, a)
                    if x:
                        _logger.debug('{0}: {1}'.format(a.upper(), x.decode('utf-8').strip()))
                raise RuntimeError('Failed to render locales successfully')
        cfg = os.path.join(self.chroot_base, 'etc', 'locale.conf')
        # And now we write the variables.
        # We have to strip out the section from the ini.
        cfgbuf = io.StringIO()
        self._localevars.write(cfgbuf, space_around_delimiters = False)
        cfgbuf.seek(0, 0)
        with open(cfg, 'w') as fh:
            for line in cfgbuf.readlines():
                if line.startswith('[BASE]') or line.strip() == '':
                    continue
                fh.write(line)
        os.chmod(cfg, 0o0644)
        os.chown(cfg, 0, 0)
        _logger.info('Wrote: {0}'.format(cfg))
        return(None)


class Timezone(object):
    def __init__(self, chroot_base, timezone):
        self.tz = timezone.strip().replace('.', '/')
        self.chroot_base = chroot_base

    def _verify(self):
        tzfilebase = os.path.join('usr', 'share', 'zoneinfo', self.tz)
        tzfile = os.path.join(self.chroot_base, tzfilebase)
        if not os.path.isfile(tzfile):
            _logger.error('Timezone {0} does not have a matching timezone file on target system.'.format(self.tz))
            raise ValueError('Invalid timezone')
        return(tzfilebase)

    def apply(self):
        tzsrcfile = os.path.join('/', self._verify())
        tzdestfile = os.path.join(self.chroot_base, 'etc', 'localtime')
        if os.path.isfile(tzdestfile):
            os.remove(tzdestfile)
        os.symlink(tzsrcfile, tzdestfile)
        _logger.info('Created symlink: {0} => {1}'.format(tzsrcfile, tzdestfile))
        return(None)