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