Manage /etc/security/limits.conf (#1036)

Add operation and fact to set and fetch values from /etc/security/limits.conf Co-authored-by: Nick Mills-Barrett <[email protected]>
pyinfra-dev · Nov 26, 2023 · 9bc79ae · 9bc79ae
1 parent 8c2751d
commit 9bc79ae
Show file tree

Hide file tree

Showing 4 changed files with 205 additions and 0 deletions.
diff --git a/pyinfra/facts/server.py b/pyinfra/facts/server.py
@@ -677,3 +677,85 @@ def process(self, output):
         # replace utf8 with UTF-8 to match names in /etc/locale.gen
         # return a list of enabled locales
         return [line.replace("utf8", "UTF-8") for line in output]
+
+
+class SecurityLimits(FactBase):
+    """
+    Returns a list of security limits on the target host.
+
+    .. code:: python
+
+        [
+            {
+                "domain": "*",
+                "limit_type": "soft",
+                "item": "nofile",
+                "value": "1048576"
+            },
+            {
+                "domain": "*",
+                "limit_type": "hard",
+                "item": "nofile",
+                "value": "1048576"
+            },
+            {
+                "domain": "root",
+                "limit_type": "soft",
+                "item": "nofile",
+                "value": "1048576"
+            },
+            {
+                "domain": "root",
+                "limit_type": "hard",
+                "item": "nofile",
+                "value": "1048576"
+            },
+            {
+                "domain": "*",
+                "limit_type": "soft",
+                "item": "memlock",
+                "value": "unlimited"
+            },
+            {
+                "domain": "*",
+                "limit_type": "hard",
+                "item": "memlock",
+                "value": "unlimited"
+            },
+            {
+                "domain": "root",
+                "limit_type": "soft",
+                "item": "memlock",
+                "value": "unlimited"
+            },
+            {
+                "domain": "root",
+                "limit_type": "hard",
+                "item": "memlock",
+                "value": "unlimited"
+            }
+        ]
+    """
+
+    command = "cat /etc/security/limits.conf"
+    default = list
+
+    def process(self, output):
+        limits = []
+
+        for line in output:
+            if line.startswith("#") or not len(line.strip()):
+                continue
+
+            domain, limit_type, item, value = line.split()
+
+            limits.append(
+                {
+                    "domain": domain,
+                    "limit_type": limit_type,
+                    "item": item,
+                    "value": value,
+                },
+            )
+
+        return limits
diff --git a/pyinfra/operations/server.py b/pyinfra/operations/server.py
@@ -1219,3 +1219,40 @@ def locale(
         )
 
         yield "locale-gen"
+
+
+@operation
+def security_limit(
+    domain,
+    limit_type,
+    item,
+    value,
+):
+    """
+    Edit /etc/security/limits.conf configuration.
+
+    + domain: the domain (user, group, or wildcard) for the limit
+    + limit_type: the type of limit (hard or soft)
+    + item: the item to limit (e.g., nofile, nproc)
+    + value: the value for the limit
+
+    **Example:**
+
+    .. code:: python
+
+        security_limit(
+            name="Set nofile limit for all users",
+            domain='*',
+            limit_type='soft',
+            item='nofile',
+            value='1024',
+        )
+    """
+
+    line_format = f"{domain}\t{limit_type}\t{item}\t{value}"
+
+    yield from files.line(
+        path="/etc/security/limits.conf",
+        line=f"^{domain}[[:space:]]+{limit_type}[[:space:]]+{item}",
+        replace=line_format,
+    )
diff --git a/tests/facts/server.SecurityLimits/security_limits.json b/tests/facts/server.SecurityLimits/security_limits.json
@@ -0,0 +1,74 @@
+{
+    "command": "cat /etc/security/limits.conf",
+    "output": [
+        "# <domain>      <type>  <item>         <value>",
+        "#*               soft    core            ",
+        "#root            hard    core            10000",
+        "#*               hard    rss             1000",
+        "#@student        hard    nproc           2",
+        "#@faculty        soft    nproc           2",
+        "#@faculty        hard    nproc           5",
+        "#ftp             hard    nproc           ",
+        "#ftp             -       chroot          /ft",
+        "#@student        -       maxlogins       ",
+        "* soft nofile 104857",
+        "* hard nofile 104857",
+        "root soft nofile 104857",
+        "root hard nofile 104857",
+        "* soft memlock unlimite",
+        "* hard memlock unlimite",
+        "root soft memlock unlimite",
+        "root hard memlock unlimite"
+    ],
+    "fact": [
+        {
+            "domain": "*",
+            "limit_type": "soft",
+            "item": "nofile",
+            "value": "104857"
+        },
+        {
+            "domain": "*",
+            "limit_type": "hard",
+            "item": "nofile",
+            "value": "104857"
+        },
+        {
+            "domain": "root",
+            "limit_type": "soft",
+            "item": "nofile",
+            "value": "104857"
+        },
+        {
+            "domain": "root",
+            "limit_type": "hard",
+            "item": "nofile",
+            "value": "104857"
+        },
+        {
+            "domain": "*",
+            "limit_type": "soft",
+            "item": "memlock",
+            "value": "unlimite"
+        },
+        {
+            "domain": "*",
+            "limit_type": "hard",
+            "item": "memlock",
+            "value": "unlimite"
+        },
+        {
+            "domain": "root",
+            "limit_type": "soft",
+            "item": "memlock",
+            "value": "unlimite"
+        },
+        {
+            "domain": "root",
+            "limit_type": "hard",
+            "item": "memlock",
+            "value": "unlimite"
+        }
+    ]
+}
+
diff --git a/tests/operations/server.security_limit/set.json b/tests/operations/server.security_limit/set.json
@@ -0,0 +1,12 @@
+{
+    "kwargs": {
+        "domain": "root",
+        "limit_type": "hard",
+        "item": "memlock",
+        "value": "unlimited"
+    },
+    "facts": {"files.FindInFile": {}},
+    "commands": [
+        "if [ -f /etc/security/limits.conf ]; then ( grep '^root[[:space:]]+hard[[:space:]]+memlock.*$' /etc/security/limits.conf && sed -i.a-timestamp 's/^root[[:space:]]+hard[[:space:]]+memlock.*$/root\thard\tmemlock\tunlimited/' /etc/security/limits.conf && rm -f /etc/security/limits.conf.a-timestamp ) 2> /dev/null || echo 'root\thard\tmemlock\tunlimited' >> /etc/security/limits.conf ; else echo 'root\thard\tmemlock\tunlimited' >> /etc/security/limits.conf ; fi"
+    ]
+}