Python yml配置文件实例

目录

安装

封装实例

目录结构

功能封装

运行结果

---

安装

pip3 install PyYaml

封装实例

目录结构

test_yml/
├── config
│???├── config.py
│???├── config.yml
│???└── __pycache__
│??? ? ?└── config.cpython-38.pyc
├── lib
│???├── pathutil.py
│???├── __pycache__
│???│???├── pathutil.cpython-38.pyc
│???│???└── singleton.cpython-38.pyc
│???└── singleton.py
└── main.py

功能封装

lib/singleton.py（单例模式）

#! /usr/bin/env python
# -*- coding: utf-8 -*-


# 单例模式
def Singleton(cls):
    _instance={}
    def _singleton(*args,**kwagrs):
        if cls not in  _instance:
            _instance[cls]=cls(*args,**kwagrs)
        return _instance[cls]
    return _singleton

lib/pathutil.py（常用获取路径封装）

#! /usr/bin/env python
# -*- coding: utf-8 -*-


import os 


def base_dir():
    return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def config_dir():
    return os.path.join(base_dir(), "config")

def path_join(path):
    return os.path.join(base_dir(), path)

config/config.py（获取配置）

#! /usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import yaml
import traceback

from lib.singleton import Singleton
from lib.pathutil import path_join


@Singleton
class Config(object):
    def __init__(self) -> None:
        self.config = None

    def __parse_config(self) -> None:  # sourcery skip: simplify-generator
        config = None
        self.config = {
            'host' : [],
            'port' : [],
            'cookie' : '',
            'mrs' : []
        }

        with open(path_join('config/config.yml'), 'r', encoding="utf-8") as f:
            config = yaml.load(f, Loader=yaml.SafeLoader)
        try:
            for h in config['host']:
                self.config['host'].append(h)
            for p in config['port']:
                self.config['port'].append(p)
            self.config['cookie'] = config['cookie'] 
            for m in config['matchreplace']:
                self.config['mrs'].append(m)
        except Exception as e:
            print("Error: 配置解析出错，请检查")
            traceback.print_exc()
            exit(0)

    def get_config(self) -> dict:
        if self.config is None:
            self.__parse_config()
        return self.config

    def check_config(self) -> bool :
        if ('host' not in self.config) or ('cookie' not in self.config) or ('mrs' not in self.config):
            print("Error: 配置出错，请检查")
            return False
        if type(self.config['host']) != list or type(self.config['mrs']) != list or type(self.config['cookie']) != str:
            print("Error: 配置出错，请检查")
            return False 
        if self.config['cookie'] == '':
            print("Error: 配置出错，请检查，cookie不能为空")
            return False
        return True

config/config.yml（yml配置文件）

# 对于哪些域名进行越权漏洞检测，可以写正则表达式，正则re.compile()不通过的会被认为是纯字符串，不含端口
host:
  - "localhost"
  - ".*\\.aaa\\.com"
  - "127.0.0.1"

# port表示对上述域名的哪些端口, 写数字就表示精准匹配端口，写字符串表示端口的正则表达式，如果不想限制端口直接写 ".*" 即可(表示匹配任意)
port:
  - 80
  - "1\\d{4}"
  - ".*"

# cookie直接全量替换，所以要写全，有多少个cookie字段写多少个。
cookie: "JSESSION=Y1ng.VIP; x-csrf-token=123123123; signature=12346576788"

# matchreplace会做替换
  # pattern可以写正则表达式，如果不能成功执行正则re.compile()则会被认为是纯字符串，，replace则为替换后的字符串（而非正则表达式）
  # location可以有以下几种选项
    # URL 对URL做替换
    # PATH 对PATH路径做替换
    # HEADER 对请求头做替换，当替换HEADER时，replace为一个对象，name表示对应头部名，value是匹配的头部值的正则
    # BODY 对POST传输的BODY部分做替换
matchreplace: 
  - {
      location: "URL",
      pattern: "_signature=2Cw00QAAAAA9GaIS4KF.M9gsNMAALwH",
      replace: "_signature=2RlsxgAAAAA8LPoFKkjFZNkZbNAALrU",
    }
  - {
      location: "HEADER",
      pattern:  "Bareer .*$" ,
      replace:  { name: "Authorization", value: "Bareer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" },
    }

main.py（测试的主文件）

#! /usr/bin/env python
# -*- coding: utf-8 -*-


from config.config import Config


if __name__ == "__main__":
    config = Config().get_config()
    if not Config().check_config():
        exit(0)
    print(config)

运行结果

{
	'host': ['localhost', '.*\\.aaa\\.com', '127.0.0.1'],
	'port': [80, '1\\d{4}', '.*'],
	'cookie': 'JSESSION=Y1ng.VIP; x-csrf-token=123123123; signature=12346576788',
	'mrs': [{
		'location': 'URL',
		'pattern': '_signature=2Cw00QAAAAA9GaIS4KF.M9gsNMAALwH',
		'replace': '_signature=2RlsxgAAAAA8LPoFKkjFZNkZbNAALrU'
	}, {
		'location': 'HEADER',
		'pattern': 'Bareer .*$',
		'replace': {
			'name': 'Authorization',
			'value': 'Bareer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c'
		}
	}]
}