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):
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                wants = systemdir / "{}.{}".format(dependent, dirstem) / service
205                add_link(wants, target)
206
207        except KeyError:
208            pass
209
210    def enable(self, caller_unit=None):
211        # if we're enabling an instance, first extract the actual instance
212        # then figure out what the template unit is
213        template = re.match(r"[^@]+@(?P<instance>[^\.]*)\.", self.unit)
214        instance_unit_name = None
215        if template:
216            instance = template.group('instance')
217            if instance != "":
218                instance_unit_name = self.unit
219            unit = re.sub(r"@[^\.]*\.", "@.", self.unit, 1)
220        else:
221            instance = None
222            unit = self.unit
223
224        path = self._path_for_unit(unit)
225
226        if path.is_symlink():
227            # ignore aliases
228            return
229
230        config = SystemdFile(self.root, path, instance_unit_name)
231        if instance == "":
232            try:
233                default_instance = config.get('Install', 'DefaultInstance')[0]
234            except KeyError:
235                # no default instance, so nothing to enable
236                return
237
238            service = self.unit.replace("@.",
239                                        "@{}.".format(default_instance))
240        else:
241            service = self.unit
242
243        self._process_deps(config, service, path, 'WantedBy', 'wants')
244        self._process_deps(config, service, path, 'RequiredBy', 'requires')
245
246        try:
247            for also in config.get('Install', 'Also'):
248                try:
249                    if caller_unit != also:
250                        SystemdUnit(self.root, also).enable(unit)
251                except SystemdUnitNotFoundError as e:
252                    sys.exit("Error: Systemctl also enable issue with  %s (%s)" % (service, e.unit))
253
254        except KeyError:
255            pass
256
257        systemdir = self.root / SYSCONFDIR / "systemd" / "system"
258        target = ROOT / path.relative_to(self.root)
259        try:
260            for dest in config.get('Install', 'Alias'):
261                alias = systemdir / dest
262                add_link(alias, target)
263
264        except KeyError:
265            pass
266
267    def mask(self):
268        systemdir = self.root / SYSCONFDIR / "systemd" / "system"
269        add_link(systemdir / self.unit, "/dev/null")
270
271
272def collect_services(root):
273    """Collect list of service files"""
274    services = set()
275    for location in locations:
276        paths = (root / location / "system").glob("*")
277        for path in paths:
278            if path.is_dir():
279                continue
280            services.add(path.name)
281
282    return services
283
284
285def preset_all(root):
286    presets = Presets('system-preset', root)
287    services = collect_services(root)
288
289    for service in services:
290        state = presets.state(service)
291
292        if state == "enable" or state is None:
293            try:
294                SystemdUnit(root, service).enable()
295            except SystemdUnitNotFoundError:
296                sys.exit("Error: Systemctl preset_all issue in %s" % service)
297
298    # If we populate the systemd links we also create /etc/machine-id, which
299    # allows systemd to boot with the filesystem read-only before generating
300    # a real value and then committing it back.
301    #
302    # For the stateless configuration, where /etc is generated at runtime
303    # (for example on a tmpfs), this script shouldn't run at all and we
304    # allow systemd to completely populate /etc.
305    (root / SYSCONFDIR / "machine-id").touch()
306
307
308def main():
309    if sys.version_info < (3, 4, 0):
310        sys.exit("Python 3.4 or greater is required")
311
312    parser = argparse.ArgumentParser()
313    parser.add_argument('command', nargs='?', choices=['enable', 'mask',
314                                                     'preset-all'])
315    parser.add_argument('service', nargs=argparse.REMAINDER)
316    parser.add_argument('--root')
317    parser.add_argument('--preset-mode',
318                        choices=['full', 'enable-only', 'disable-only'],
319                        default='full')
320
321    args = parser.parse_args()
322
323    root = Path(args.root) if args.root else ROOT
324
325    locations.append(SYSCONFDIR / "systemd")
326    # Handle the usrmerge case by ignoring /lib when it's a symlink
327    if not (root / BASE_LIBDIR).is_symlink():
328        locations.append(BASE_LIBDIR / "systemd")
329    locations.append(LIBDIR / "systemd")
330
331    command = args.command
332    if not command:
333        parser.print_help()
334        return 0
335
336    if command == "mask":
337        for service in args.service:
338            try:
339                SystemdUnit(root, service).mask()
340            except SystemdUnitNotFoundError as e:
341                sys.exit("Error: Systemctl main mask issue in %s (%s)" % (service, e.unit))
342    elif command == "enable":
343        for service in args.service:
344            try:
345                SystemdUnit(root, service).enable()
346            except SystemdUnitNotFoundError as e:
347                sys.exit("Error: Systemctl main enable issue in %s (%s)" % (service, e.unit))
348    elif command == "preset-all":
349        if len(args.service) != 0:
350            sys.exit("Too many arguments.")
351        if args.preset_mode != "enable-only":
352            sys.exit("Only enable-only is supported as preset-mode.")
353        preset_all(root)
354    else:
355        raise RuntimeError()
356
357
358if __name__ == '__main__':
359    main()
360