1#!/usr/bin/env python3
2"""systemctl: subset of systemctl used for image construction
3
4Mask/preset systemd units
5"""
6
7import argparse
8import fnmatch
9import os
10import re
11import sys
12
13from collections import namedtuple
14from itertools import chain
15from pathlib import Path
16
17version = 1.0
18
19ROOT = Path("/")
20SYSCONFDIR = Path("etc")
21BASE_LIBDIR = Path("lib")
22LIBDIR = Path("usr", "lib")
23
24locations = list()
25
26
27class SystemdFile():
28    """Class representing a single systemd configuration file"""
29
30    _clearable_keys = ['WantedBy']
31
32    def __init__(self, root, path, instance_unit_name):
33        self.sections = dict()
34        self._parse(root, path)
35        dirname = os.path.basename(path.name) + ".d"
36        for location in locations:
37            files = (root / location / "system" / dirname).glob("*.conf")
38            if instance_unit_name:
39                inst_dirname = instance_unit_name + ".d"
40                files = chain(files, (root / location / "system" / inst_dirname).glob("*.conf"))
41            for path2 in sorted(files):
42                self._parse(root, path2)
43
44    def _parse(self, root, path):
45        """Parse a systemd syntax configuration file
46
47        Args:
48            path: A pathlib.Path object pointing to the file
49
50        """
51        skip_re = re.compile(r"^\s*([#;]|$)")
52        section_re = re.compile(r"^\s*\[(?P<section>.*)\]")
53        kv_re = re.compile(r"^\s*(?P<key>[^\s]+)\s*=\s*(?P<value>.*)")
54        section = None
55
56        if path.is_symlink():
57            try:
58                path.resolve()
59            except FileNotFoundError:
60                # broken symlink, try relative to root
61                path = root / Path(os.readlink(str(path))).relative_to(ROOT)
62
63        with path.open() as f:
64            for line in f:
65                if skip_re.match(line):
66                    continue
67
68                line = line.strip()
69                m = section_re.match(line)
70                if m:
71                    if m.group('section') not in self.sections:
72                        section = dict()
73                        self.sections[m.group('section')] = section
74                    else:
75                        section = self.sections[m.group('section')]
76                    continue
77
78                while line.endswith("\\"):
79                    line += f.readline().rstrip("\n")
80
81                m = kv_re.match(line)
82                k = m.group('key')
83                v = m.group('value')
84                if k not in section:
85                    section[k] = list()
86
87                # If we come across a "key=" line for a "clearable key", then
88                # forget all preceding assignments. This works because we are
89                # processing files in correct parse order.
90                if k in self._clearable_keys and not v:
91                    del section[k]
92                    continue
93
94                section[k].extend(v.split())
95
96    def get(self, section, prop):
97        """Get a property from section
98
99        Args:
100            section: Section to retrieve property from
101            prop: Property to retrieve
102
103        Returns:
104            List representing all properties of type prop in section.
105
106        Raises:
107            KeyError: if ``section`` or ``prop`` not found
108        """
109        return self.sections[section][prop]
110
111
112class Presets():
113    """Class representing all systemd presets"""
114    def __init__(self, scope, root):
115        self.directives = list()
116        self._collect_presets(scope, root)
117
118    def _parse_presets(self, presets):
119        """Parse presets out of a set of preset files"""
120        skip_re = re.compile(r"^\s*([#;]|$)")
121        directive_re = re.compile(r"^\s*(?P<action>enable|disable)\s+(?P<unit_name>(.+))")
122
123        Directive = namedtuple("Directive", "action unit_name")
124        for preset in presets:
125            with preset.open() as f:
126                for line in f:
127                    m = directive_re.match(line)
128                    if m:
129                        directive = Directive(action=m.group('action'),
130                                              unit_name=m.group('unit_name'))
131                        self.directives.append(directive)
132                    elif skip_re.match(line):
133                        pass
134                    else:
135                        sys.exit("Unparsed preset line in {}".format(preset))
136
137    def _collect_presets(self, scope, root):
138        """Collect list of preset files"""
139        presets = dict()
140        for location in locations:
141            paths = (root / location / scope).glob("*.preset")
142            for path in paths:
143                # earlier names override later ones
144                if path.name not in presets:
145                    presets[path.name] = path
146
147        self._parse_presets([v for k, v in sorted(presets.items())])
148
149    def state(self, unit_name):
150        """Return state of preset for unit_name
151
152        Args:
153            presets: set of presets
154            unit_name: name of the unit
155
156        Returns:
157            None: no matching preset
158            `enable`: unit_name is enabled
159            `disable`: unit_name is disabled
160        """
161        for directive in self.directives:
162            if fnmatch.fnmatch(unit_name, directive.unit_name):
163                return directive.action
164
165        return None
166
167
168def add_link(path, target):
169    try:
170        path.parent.mkdir(parents=True)
171    except FileExistsError:
172        pass
173    if not path.is_symlink():
174        print("ln -s {} {}".format(target, path))
175        path.symlink_to(target)
176
177
178class SystemdUnitNotFoundError(Exception):
179    def __init__(self, path, unit):
180        self.path = path
181        self.unit = unit
182
183
184class SystemdUnit():
185    def __init__(self, root, unit):
186        self.root = root
187        self.unit = unit
188        self.config = None
189
190    def _path_for_unit(self, unit):
191        for location in locations:
192            path = self.root / location / "system" / unit
193            if path.exists() or path.is_symlink():
194                return path
195
196        raise SystemdUnitNotFoundError(self.root, unit)
197
198    def _process_deps(self, config, service, location, prop, dirstem, instance):
199        systemdir = self.root / SYSCONFDIR / "systemd" / "system"
200
201        target = ROOT / location.relative_to(self.root)
202        try:
203            for dependent in config.get('Install', prop):
204                # expand any %i to instance (ignoring escape sequence %%)
205                dependent = re.sub("([^%](%%)*)%i", "\\g<1>{}".format(instance), dependent)
206                wants = systemdir / "{}.{}".format(dependent, dirstem) / service
207                add_link(wants, target)
208
209        except KeyError:
210            pass
211
212    def enable(self, units_enabled=[]):
213        # if we're enabling an instance, first extract the actual instance
214        # then figure out what the template unit is
215        template = re.match(r"[^@]+@(?P<instance>[^\.]*)\.", self.unit)
216        instance_unit_name = None
217        if template:
218            instance = template.group('instance')
219            if instance != "":
220                instance_unit_name = self.unit
221            unit = re.sub(r"@[^\.]*\.", "@.", self.unit, 1)
222        else:
223            instance = None
224            unit = self.unit
225
226        path = self._path_for_unit(unit)
227
228        if path.is_symlink():
229            # ignore aliases
230            return
231
232        config = SystemdFile(self.root, path, instance_unit_name)
233        if instance == "":
234            try:
235                default_instance = config.get('Install', 'DefaultInstance')[0]
236            except KeyError:
237                # no default instance, so nothing to enable
238                return
239
240            service = self.unit.replace("@.",
241                                        "@{}.".format(default_instance))
242        else:
243            service = self.unit
244
245        self._process_deps(config, service, path, 'WantedBy', 'wants', instance)
246        self._process_deps(config, service, path, 'RequiredBy', 'requires', instance)
247
248        try:
249            for also in config.get('Install', 'Also'):
250                try:
251                     units_enabled.append(unit)
252                     if also not in units_enabled:
253                        SystemdUnit(self.root, also).enable(units_enabled)
254                except SystemdUnitNotFoundError as e:
255                    sys.exit("Error: Systemctl also enable issue with  %s (%s)" % (service, e.unit))
256
257        except KeyError:
258            pass
259
260        systemdir = self.root / SYSCONFDIR / "systemd" / "system"
261        target = ROOT / path.relative_to(self.root)
262        try:
263            for dest in config.get('Install', 'Alias'):
264                alias = systemdir / dest
265                add_link(alias, target)
266
267        except KeyError:
268            pass
269
270    def mask(self):
271        systemdir = self.root / SYSCONFDIR / "systemd" / "system"
272        add_link(systemdir / self.unit, "/dev/null")
273
274
275def collect_services(root):
276    """Collect list of service files"""
277    services = set()
278    for location in locations:
279        paths = (root / location / "system").glob("*")
280        for path in paths:
281            if path.is_dir():
282                continue
283            services.add(path.name)
284
285    return services
286
287
288def preset_all(root):
289    presets = Presets('system-preset', root)
290    services = collect_services(root)
291
292    for service in services:
293        state = presets.state(service)
294
295        if state == "enable" or state is None:
296            try:
297                SystemdUnit(root, service).enable()
298            except SystemdUnitNotFoundError:
299                sys.exit("Error: Systemctl preset_all issue in %s" % service)
300
301    # If we populate the systemd links we also create /etc/machine-id, which
302    # allows systemd to boot with the filesystem read-only before generating
303    # a real value and then committing it back.
304    #
305    # For the stateless configuration, where /etc is generated at runtime
306    # (for example on a tmpfs), this script shouldn't run at all and we
307    # allow systemd to completely populate /etc.
308    (root / SYSCONFDIR / "machine-id").touch()
309
310
311def main():
312    if sys.version_info < (3, 4, 0):
313        sys.exit("Python 3.4 or greater is required")
314
315    parser = argparse.ArgumentParser()
316    parser.add_argument('command', nargs='?', choices=['enable', 'mask',
317                                                     'preset-all'])
318    parser.add_argument('service', nargs=argparse.REMAINDER)
319    parser.add_argument('--root')
320    parser.add_argument('--preset-mode',
321                        choices=['full', 'enable-only', 'disable-only'],
322                        default='full')
323
324    args = parser.parse_args()
325
326    root = Path(args.root) if args.root else ROOT
327
328    locations.append(SYSCONFDIR / "systemd")
329    # Handle the usrmerge case by ignoring /lib when it's a symlink
330    if not (root / BASE_LIBDIR).is_symlink():
331        locations.append(BASE_LIBDIR / "systemd")
332    locations.append(LIBDIR / "systemd")
333
334    command = args.command
335    if not command:
336        parser.print_help()
337        return 0
338
339    if command == "mask":
340        for service in args.service:
341            try:
342                SystemdUnit(root, service).mask()
343            except SystemdUnitNotFoundError as e:
344                sys.exit("Error: Systemctl main mask issue in %s (%s)" % (service, e.unit))
345    elif command == "enable":
346        for service in args.service:
347            try:
348                SystemdUnit(root, service).enable()
349            except SystemdUnitNotFoundError as e:
350                sys.exit("Error: Systemctl main enable issue in %s (%s)" % (service, e.unit))
351    elif command == "preset-all":
352        if len(args.service) != 0:
353            sys.exit("Too many arguments.")
354        if args.preset_mode != "enable-only":
355            sys.exit("Only enable-only is supported as preset-mode.")
356        preset_all(root)
357    else:
358        raise RuntimeError()
359
360
361if __name__ == '__main__':
362    main()
363