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