Skip to content

RESULTS 类型插件

RESULTS 类型插件用来自定义检测结果的导出,例如导出 html 报表等。

from pocsuite3.api import PluginBase
from pocsuite3.api import PLUGIN_TYPE
from pocsuite3.api import logger
from pocsuite3.api import get_results
from pocsuite3.api import register_plugin

class HtmlReport(PluginBase):
    category = PLUGIN_TYPE.RESULTS

    def init(self):
        debug_msg = "[PLUGIN] html_report plugin init..."
        logger.debug(debug_msg)

    def start(self):
        # TODO
        # Generate html report

        for result in get_results():
            pass

        info_msg = '[PLUGIN] generate html report done.'
        logger.info(info_msg)

register_plugin(HtmlReport)

若需要实时的保存结果,需要申明 handle 来处理,代码示例如下:

import os
import time
import json
from pocsuite3.api import PLUGIN_TYPE
from pocsuite3.api import PluginBase
from pocsuite3.api import logger
from pocsuite3.api import register_plugin, paths, conf


class FileRecord(PluginBase):
    category = PLUGIN_TYPE.RESULTS
    filename = conf.output_path or os.path.join(paths.POCSUITE_OUTPUT_PATH, "{}.txt".format(int(time.time())))
    file = None

    def init(self):
        debug_msg = "[PLUGIN] file_record plugin init..."
        logger.debug(debug_msg)
        logger.info("[PLUGIN] The result will be recorded in {}".format(self.filename))
        if os.path.exists(self.filename):
            raise Exception("The {} has existed".format(self.filename))
        self.file = open(self.filename, 'a+')

    def handle(self, output):
        status = output.get("status")
        if status and status == "success":
            poc_name = output.get("poc_name")
            target = output.get("target")
            created = output.get("created")
            result = output.get("result")
            msg = {
                "target": target,
                "poc_name": poc_name,
                "result": result,
                "created_time": created
            }
            self.file.write(json.dumps(msg) + '\n')

    def start(self):
        self.file.close()
        msg = "[PLUGIN] File saved in {}".format(self.filename)
        logger.info(msg)


register_plugin(FileRecord)

Released under the GPLv2 License.