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", "\\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, caller_unit=None):
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                    if caller_unit != also:
252                        SystemdUnit(self.root, also).enable(unit)
253                except SystemdUnitNotFoundError as e:
254                    sys.exit("Error: Systemctl also enable issue with  %s (%s)" % (service, e.unit))
255
256        except KeyError:
257            pass
258
259        systemdir = self.root / SYSCONFDIR / "systemd" / "system"
260        target = ROOT / path.relative_to(self.root)
261        try:
262            for dest in config.get('Install', 'Alias'):
263                alias = systemdir / dest
264                add_link(alias, target)
265
266        except KeyError:
267            pass
268
269    def mask(self):
270        systemdir = self.root / SYSCONFDIR / "systemd" / "system"
271        add_link(systemdir / self.unit, "/dev/null")
272
273
274def collect_services(root):
275    """Collect list of service files"""
276    services = set()
277    for location in locations:
278        paths = (root / location / "system").glob("*")
279        for path in paths:
280            if path.is_dir():
281                continue
282            services.add(path.name)
283
284    return services
285
286
287def preset_all(root):
288    presets = Presets('system-preset', root)
289    services = collect_services(root)
290
291    for service in services:
292        state = presets.state(service)
293
294        if state == "enable" or state is None:
295            try:
296                SystemdUnit(root, service).enable()
297            except SystemdUnitNotFoundError:
298                sys.exit("Error: Systemctl preset_all issue in %s" % service)
299
300    # If we populate the systemd links we also create /etc/machine-id, which
301    # allows systemd to boot with the filesystem read-only before generating
302    # a real value and then committing it back.
303    #
304    # For the stateless configuration, where /etc is generated at runtime
305    # (for example on a tmpfs), this script shouldn't run at all and we
306    # allow systemd to completely populate /etc.
307    (root / SYSCONFDIR / "machine-id").touch()
308
309
310def main():
311    if sys.version_info < (3, 4, 0):
312        sys.exit("Python 3.4 or greater is required")
313
314    parser = argparse.ArgumentParser()
315    parser.add_argument('command', nargs='?', choices=['enable', 'mask',
316                                                     'preset-all'])
317    parser.add_argument('service', nargs=argparse.REMAINDER)
318    parser.add_argument('--root')
319    parser.add_argument('--preset-mode',
320                        choices=['full', 'enable-only', 'disable-only'],
321                        default='full')
322
323    args = parser.parse_args()
324
325    root = Path(args.root) if args.root else ROOT
326
327    locations.append(SYSCONFDIR / "systemd")
328    # Handle the usrmerge case by ignoring /lib when it's a symlink
329    if not (root / BASE_LIBDIR).is_symlink():
330        locations.append(BASE_LIBDIR / "systemd")
331    locations.append(LIBDIR / "systemd")
332
333    command = args.command
334    if not command:
335        parser.print_help()
336        return 0
337
338    if command == "mask":
339        for service in args.service:
340            try:
341                SystemdUnit(root, service).mask()
342            except SystemdUnitNotFoundError as e:
343                sys.exit("Error: Systemctl main mask issue in %s (%s)" % (service, e.unit))
344    elif command == "enable":
345        for service in args.service:
346            try:
347                SystemdUnit(root, service).enable()
348            except SystemdUnitNotFoundError as e:
349                sys.exit("Error: Systemctl main enable issue in %s (%s)" % (service, e.unit))
350    elif command == "preset-all":
351        if len(args.service) != 0:
352            sys.exit("Too many arguments.")
353        if args.preset_mode != "enable-only":
354            sys.exit("Only enable-only is supported as preset-mode.")
355        preset_all(root)
356    else:
357        raise RuntimeError()
358
359
360if __name__ == '__main__':
361    main()
362