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