只允许python脚本以指定的方式运行

13. 只允许python脚本以指定的方式运行

​ 有时候我们想限制脚本的运行方法, 比如我们有test.py 这个脚本, 可以直接使用 python test.py 运行, 而 /home/test/python3.8/bin/python 是 python的绝对路径, 也可以使用 /home/test/python3.8/bin/python test.py 来运行该脚本

​ 现在我们想 test.py 这个脚本只能用 /home/test/python3.8/bin/python test.py 这种方式来运行,

思路: 使用psutil 获取test.py 脚本的启动命令, 然后和指定的启动方法进行比对

# -*- encoding: utf-8 -*-

import re
import sys
import time
import psutil


def get_run_cmd():

    """
        功能: 找出当前进程的启动命令
        参数: 
            None
        返回值:
            cmdline: 字符串, 启动当前进程的命令
            None: 没有找到该进程的启动命令
    """

    # 拼接命令行参数, sys.argv 只能获取脚本和参数
    # 无法完整获取启动命令
    # 有个问题, 如果我们指定命令的参数有用空格
    # 如 python test.py "hello world"
    # 通过这个方式拼接, 双引号或者单引号会被去掉变成, 如下就是拼接好的字符串
    # python test.py hello world
    sub_run_cmd = " ".join(sys.argv)

    # 遍历系统上所有的进程
    for process in psutil.process_iter():   

        try:
            
            # 拼接进程的执行命令
            cmdline = " ".join(process.cmdline())

            # 由于获取到的参数列表是完整参数的一部分, 
            # 可以使用正则表达式进行匹配
            pattern = f"{sub_run_cmd}$"

            # 匹配上
            if re.search(pattern=pattern, string=cmdline):

                return cmdline
        
        # 没有权限, 跳过该进程
        except psutil.AccessDenied:

            pass
        
        # 进程不存在或者已经执行完了, 跳过
        except psutil.NoSuchProcess:
            
            pass

    return None

if __name__ == "__main__":


    run_cmd = get_run_cmd()
    
    # 指定的进程启动命令
    start_cmd_appointed = "/home/test/python3.8/bin/python -u test.py \"1 2\""
	
	# 由于上面的拼接会去掉单引号和双引号, 所有我们这里比对也要去掉
    new_start_cmd_appointed = start_cmd_appointed.replace("\"", "").replace("\'", "")

    if run_cmd == new_start_cmd_appointed:

        print("启动成功")
    
    else:

        print(f"必须用 {start_cmd_appointed} 命令来执行")

        exit(1)