【第三方插件】处理输出结果

[charis3306]

处理输出结果

对-o输出结果进行处理

#时间：2023/08/14
#功能：dismap 数据处理工具
#作者：charis ks安全团队
#参数：-f 要处理的文件位置 -o 输出文件位置

import re
from openpyxl import Workbook
from openpyxl.styles import PatternFill, Border, Side, Alignment, Font
from typing import Tuple, Optional
import argparse
from openpyxl.worksheet.worksheet import Worksheet


class DismapDataProcessor:
    def __init__(self, input_path: str, output_path: str):
        self.input_path = input_path
        self.output_path = output_path
        self.wb: Workbook = Workbook()
        self.ws: Worksheet = self.wb.active
        self.ws.title = "dismap_data"

    def create_excel_workbook(self) -> None:
        # 添加列名并设置样式
        column_names = ["IP地址", "端口", "协议", "链接"]
        self.ws.append(column_names)

        for col in self.ws.iter_cols(min_col=1, max_col=4):
            for cell in col:
                cell.fill = PatternFill(start_color="00FF00", end_color="00FF00", fill_type="solid")
                cell.font = Font(color="000000")
                cell.alignment = Alignment(horizontal="center")

        # 设置所有单元格的边框线
        border = Border(left=Side(border_style="thin"),
                        right=Side(border_style="thin"),
                        top=Side(border_style="thin"),
                        bottom=Side(border_style="thin"))

        for row in self.ws.iter_rows(min_row=1, max_row=self.ws.max_row, min_col=1, max_col=self.ws.max_column):
            for cell in row:
                cell.border = border

    def extract_data(self, line: str) -> Optional[Tuple[str, str, str, str]]:
        try:
            # 定义正则表达式模式来提取IP地址、端口和协议
            pattern = r'((?:\w+)://([\d.]+):(\d+)|(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?):(\d+))'

            matches = re.findall(pattern, line)
            if matches:
                for match in matches:
                    # 处理其他协议
                    protocol = f'{match[0]}://{match[1]}:{match[2]}'
                    if not "://" in match[0]:
                        ip_address = re.match(
                            r"(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)",
                            match[0].strip('\n')).group()
                        return (ip_address, match[3], "tcp", f'tcp://{match[0]}')
                    else:
                        xy = re.search(r'(\w+)://', match[0]).group(1)
                        return (match[1], match[2], xy, match[0])
        except Exception as e:
            print(f"提取数据时发生错误：{str(e)}")
        return None

    def process_data(self) -> None:
        try:
            self.create_excel_workbook()

            with open(self.input_path, 'r', encoding='utf-8') as f:
                listData = f.readlines()

            for line in listData:
                data = self.extract_data(line)
                if data:
                    self.ws.append(data)

            # 保存Excel文件
            self.wb.save(self.output_path)
            print(f'数据已保存到 {self.output_path}')
        except Exception as e:
            print(f"处理文件时发生错误：{str(e)}")


def print_banner():
    banner = """
    *****************************************************
            *             dismap data         *
    *****************************************************
    """
    print(banner)


def main():
    print_banner()

    parser = argparse.ArgumentParser(description='dismap 数据处理工具')
    parser.add_argument('-f', '--input-file', required=True, type=str, help='要处理的文件路径')
    parser.add_argument('-o', '--output-file', required=True, type=str, help='要保存的Excel文件路径')

    args = parser.parse_args()

    processor = DismapDataProcessor(args.input_file, args.output_file)
    processor.process_data()


if __name__ == '__main__':
    main()

效果如下