From 364c5b8ac1aab70ebc04f7bd0d59b61b12f68085 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=AE=89=E4=BA=91=E7=9A=93?= <anyunhao@kylinos.cn>
Date: Mon, 9 Dec 2024 14:10:45 +0800
Subject: [PATCH] Add process validation for the kyclassifier command

---
 README.md                 |  2 +-
 main.py                   |  5 ++++-
 src/utils/programcheck.py | 17 +++++++++++++++++
 3 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 src/utils/programcheck.py

diff --git a/README.md b/README.md
index 1bfc4dc..46854ac 100644
--- a/README.md
+++ b/README.md
@@ -21,7 +21,7 @@ rpm -ivh ~/rpmbuild/RPMS/xxxx/kyclassifier-1.1-0.xxxx.rpm
 
 依赖安装：
 yum install python3-hawkey
-pip3 install pycdlib Jinja2 MarkupSafe
+pip3 install pycdlib Jinja2 MarkupSafe psutil
 
 
 #### 使用说明
diff --git a/main.py b/main.py
index 7df2419..fca544a 100644
--- a/main.py
+++ b/main.py
@@ -26,13 +26,13 @@ from src.utils import depparse
 from src.utils import util
 from src.utils.util import ISOUtils
 from src.utils.config import BaseConfig
+from src.utils.programcheck import find_processes_with_cmdline_keyword
 from src.main.alglayer import AlgLayer
 from src.main.algclassify import AlgClassify
 from src.log.logger import logger, LOGGER
 from src.report.report_generator import ReportGenerator
 from src.rpmquery.querylayeriniso import QueryLayerInIso
 
-
 class kyClassifier(object):
 
     @classmethod
@@ -119,6 +119,9 @@ if __name__ == '__main__':
 
     if args.console_log:
         LOGGER.update_console_log(logger)
+    
+    #进程检查函数
+    find_processes_with_cmdline_keyword('kyclassifier')
 
     logger.info("Start run kyclassifier...")
     if args.iso:
diff --git a/src/utils/programcheck.py b/src/utils/programcheck.py
new file mode 100644
index 0000000..44e74d0
--- /dev/null
+++ b/src/utils/programcheck.py
@@ -0,0 +1,17 @@
+import psutil
+import sys
+import os
+
+from src.log.logger import logger
+
+def find_processes_with_cmdline_keyword(keyword):
+
+    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
+        try:
+            cmdline = proc.info['cmdline']
+            if keyword in ' '.join(cmdline).lower() and proc.info['pid'] != os.getpid():
+                #print(f"PID: {proc.info['pid']}, Name: {proc.info['name']}, Cmdline: {' '.join(cmdline)}")
+                logger.error("Another kyclassifier process is running, the pid = %s " %proc.info['pid'])
+                sys.exit(-1)
+        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess) as e:
+            logger.error("Encountered an exception for PID {proc.info.get('pid', 'unknown')}: {e}")
\ No newline at end of file
-- 
Gitee