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