From 80a9dda7f8e2309b933d9c3f7a521bd0a0f8f3d5 Mon Sep 17 00:00:00 2001
From: Oren de Lame <oren.delame@granulate.io>
Date: Tue, 20 Aug 2024 15:35:28 +0300
Subject: [PATCH] Revert wrongfully deleted files (#258)

From https://github.com/Granulate/granulate-utils/pull/257
---
 granulate_utils/linux/containers.py | 42 +++++++++++++++++++++++++++++
 granulate_utils/linux/ns.py         | 33 ++++++++++++++++++++++-
 granulate_utils/linux/process.py    | 30 +++++++++++++++++++++
 3 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 granulate_utils/linux/containers.py

diff --git a/granulate_utils/linux/containers.py b/granulate_utils/linux/containers.py
new file mode 100644
index 00000000..0b892e19
--- /dev/null
+++ b/granulate_utils/linux/containers.py
@@ -0,0 +1,42 @@
+#
+# Copyright (C) 2023 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import re
+from typing import Optional
+
+from psutil import Process
+
+from granulate_utils.linux.process import get_process_cgroups
+
+# ECS uses /ecs/uuid/container-id
+# standard Docker uses /docker/container-id
+# k8s uses /kubepods/{burstable,besteffort}/uuid/container-id
+# there are some variations to the above formats, but generally, the container
+# ID is always 64-hex digits.
+CONTAINER_ID_PATTERN = re.compile(r"[a-f0-9]{64}")
+
+
+def get_process_container_id(process: Process) -> Optional[str]:
+    """
+    Gets the container ID of a running process, or None if not in a container.
+    :raises NoSuchProcess: If the process doesn't or no longer exists
+    """
+    for proc_cgroup_line in get_process_cgroups(process):
+        found = CONTAINER_ID_PATTERN.findall(proc_cgroup_line.relative_path)
+        if found:
+            return found[-1]
+
+    return None
diff --git a/granulate_utils/linux/ns.py b/granulate_utils/linux/ns.py
index b34bb59b..1be5a970 100644
--- a/granulate_utils/linux/ns.py
+++ b/granulate_utils/linux/ns.py
@@ -23,9 +23,10 @@
 from threading import Thread
 from typing import Callable, Collection, List, Optional, TypeVar, Union
 
-from psutil import NoSuchProcess, Process
+from psutil import NoSuchProcess, Process, process_iter
 
 from granulate_utils.exceptions import UnsupportedNamespaceError
+from granulate_utils.linux import containers
 
 T = TypeVar("T")
 
@@ -332,3 +333,33 @@ def resolve_host_path(process: Process, ns_path: str, from_ancestor: bool = True
     Get a path in the host mount namespace pointing to path in process mount namespace.
     """
     return resolve_proc_root_links(get_proc_root_path(process, from_ancestor=from_ancestor), ns_path)
+
+
+def get_host_pid(nspid: int, container_id: str) -> Optional[int]:
+    assert len(container_id) == 64, f"Invalid container id {container_id!r}"
+
+    pid_namespace = ""
+    running_processes = list(process_iter())
+
+    # Get the pid namespace of the given container
+    for process in running_processes:
+        try:
+            if container_id == containers.get_process_container_id(process):
+                pid_namespace = os.readlink(f"/proc/{process.pid}/ns/pid")
+                break
+        except (FileNotFoundError, NoSuchProcess):
+            continue
+
+    if not pid_namespace:
+        return None
+
+    # Here we need to find the process by comparing pid namespaces and not by comparing cgroups, because
+    # technically a process that's running in a container pid namespace doesn't have to share its cgroup
+    for process in running_processes:
+        try:
+            if os.readlink(f"/proc/{process.pid}/ns/pid") == pid_namespace and get_process_nspid(process) == nspid:
+                return process.pid
+        except (FileNotFoundError, NoSuchProcess, PermissionError):
+            continue
+
+    return None
diff --git a/granulate_utils/linux/process.py b/granulate_utils/linux/process.py
index c5fe1aab..59e7e980 100644
--- a/granulate_utils/linux/process.py
+++ b/granulate_utils/linux/process.py
@@ -171,3 +171,33 @@ def search_for_process(filter: Callable[[psutil.Process], bool]) -> Iterator[psu
         with contextlib.suppress(NoSuchProcess, AccessDenied):
             if is_process_running(proc) and filter(proc):
                 yield proc
+
+
+class ProcCgroupLine:
+    """
+    The format of the line:  hierarchy-ID:controller-list:relative-path
+    Example line: 1:cpu:/custom_cgroup
+
+    relative-path - the path of the cgroup the process belongs to, relative to the hierarchy mount point
+    e.g. /sys/fs/cgroup/memory on v1 or just the cgroups v2 mount on v2 e.g /sys/fs/cgroup.
+    """
+
+    hier_id: str
+    controllers: List[str]
+    relative_path: str
+
+    def __init__(self, procfs_line: str):
+        hier_id, controller_list, relative_path = procfs_line.split(":", maxsplit=2)
+        self.hier_id = hier_id
+        self.controllers = controller_list.split(",")
+        self.relative_path = relative_path
+
+
+def get_process_cgroups(process: Optional[psutil.Process] = None) -> List[ProcCgroupLine]:
+    """
+    Get the cgroups of a process in [(hier id., controllers, path)] parsed form.
+    If process is None, gets the cgroups of the current process.
+    """
+    process = process or psutil.Process()
+    text = read_proc_file(process, "cgroup").decode()
+    return [ProcCgroupLine(line) for line in text.splitlines()]