Python：自动修改 SystemD service 文件中值的脚本

python ensure_timeout_start.py /etc/systemd/system/mynodejs.service

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ensure_timeout_start.py

本作品根据 Creative Commons "CC0 1.0 Universal" (CC0 1.0) 公共领域 dedication 奉献给公共领域。

Author: Uli Köhler

SPDX-License-Identifier: CC0-1.0

确保 systemd unit 文件的 [Service] 节中有 TimeoutStartSec=600。
- 如果 [Service] 不存在，将被添加。
- 如果 TimeoutStartSec 存在（任意值），将被设置为 600。
- 如果 TimeoutStartSec 不存在，将被追加到 [Service]。

用法：
  python ensure_timeout_start.py my1.service my2.service
  python ensure_timeout_start.py /etc/systemd/system/*.service --dry-run
"""

import argparse
import os
import shutil
import sys
from configupdater import ConfigUpdater


def ensure_timeout_start(filepath: str, value: str = "600", dry_run: bool = False, backup_ext: str = ".bak") -> bool:
    """
    如果需要修改（且已修改或在 dry-run 中将会修改）则返回 True，如果已符合要求则返回 False。
    """
    updater = ConfigUpdater()
    try:
        updater.read(filepath, encoding="utf-8")
    except FileNotFoundError:
        print(f"ERROR: File not found: {filepath}", file=sys.stderr)
        return False

    changed = False

    # 确保 [Service] 节存在
    if not updater.has_section("Service"):
        updater.add_section("Service")
        changed = True

    service = updater["Service"]

    if service.has_option("TimeoutStartSec"):
        current_val = service.get("TimeoutStartSec").value.strip()
        if current_val != value:
            service["TimeoutStartSec"].value = value
            changed = True
    else:
        service["TimeoutStartSec"] = value
        changed = True

    if dry_run:
        return changed

    if changed:
        # 覆盖前创建备份
        if backup_ext:
            backup_path = filepath + backup_ext
            shutil.copy2(filepath, backup_path)

        tmp_path = filepath + ".tmp"
        with open(tmp_path, "w", encoding="utf-8", newline="") as f:
            updater.write(f)

        try:
            st = os.stat(filepath)
            os.chmod(tmp_path, st.st_mode)
            try:
                os.chown(tmp_path, st.st_uid, st.st_gid)  # 可能需要 root 权限
            except PermissionError:
                pass
        except FileNotFoundError:
            pass

        os.replace(tmp_path, filepath)

    return changed


def main():
    parser = argparse.ArgumentParser(
        description="确保一个或多个 systemd unit 文件的 [Service] 节中有 TimeoutStartSec=600。"
    )
    parser.add_argument("unit_files", nargs="+", help="unit 文件路径")
    parser.add_argument("--dry-run", action="store_true", help="只显示是否需要修改，不实际写入")
    parser.add_argument("--no-backup", action="store_true", help="不创建备份文件")
    parser.add_argument("--value", default="600", help="TimeoutStartSec 的值（默认：600）")

    args = parser.parse_args()
    backup_ext = "" if args.no_backup else ".bak"

    exit_code = 0
    for path in args.unit_files:
        changed = ensure_timeout_start(path, value=args.value, dry_run=args.dry_run, backup_ext=backup_ext)
        if args.dry_run:
            print(f"{path}: {'CHANGE NEEDED' if changed else 'NO CHANGE NEEDED'}")
        else:
            print(f"{path}: {'UPDATED' if changed else 'ALREADY OK'}")
        if not os.path.exists(path):
            exit_code = 2

    sys.exit(exit_code)


if __name__ == "__main__":
    main()