denylist.plugin

from __future__ import print_function

import logging

import empire.server.common.helpers as helpers
from empire.server.common.plugins import Plugin
from empire.server.core.agent_task_service import AgentTaskService
from empire.server.core.db import models
from empire.server.core.hooks import hooks
from empire.server.core.plugin_service import PluginService

log = logging.getLogger(__name__)

"""
The purpose of this plugin is to block certain IP addresses from connecting to the server.
It is to showcase the event-driven nature of the hook system.
"""
class Plugin(Plugin):
    def onLoad(self):
        self.info = {
            'Name': 'denylist',
            "Authors": [
                {
                    "Name": "Vincent Rose",
                    "Handle": "@Vinnybod",
                    "Link": "https://twitter.com/_vinnybod",
                },
            ],
            'Description': """
            The purpose of this plugin is to block certain IP addresses from connecting to the server.
            It is to showcase the event-driven nature of the hook system.
            """,
            'Software': '',
            'Techniques': [],
            'Comments': []
        }

        self.options = {
            'Addresses': {
                'Description': 'List of IP addresses to block. Comma Separated.',
                'Required': False,
                'Value': '127.0.0.1',
            },
        }

    def execute(self, command):
        try:
            self.ip = command['Addresses']
            return f'[*] Set IP Addresses to {self.ip}'
        except Exception as e:
            log.error(e)
            self.plugin_service.plugin_socketio_message(self.info["Name"], f"[!] {e}")
            return False

    def register(self, mainMenu):
        """
        Register hooks for the plugin
        """
        self.installPath = mainMenu.installPath
        self.main_menu = mainMenu
        self.plugin_service: PluginService = mainMenu.pluginsv2
        self.agent_task_service: AgentTaskService = mainMenu.agenttasksv2
        hooks.register_hook(hooks.AFTER_AGENT_CHECKIN_HOOK, "denylist", self.run_after_agent_checkin)

    def run_after_agent_checkin(self, Session, agent: models.Agent):
        """
        Check if agent is in the list of addresses to block. If it is, kill the agent.
        """
        if not agent.internal_ip or not self.ip:
            return

        agent_ips = agent.internal_ip.split(',')
        deny_list = self.ip.split(',')
        for ip in agent_ips:
            if ip in deny_list:
                self.plugin_service.plugin_socketio_message(self.info['Name'], f'[!] Killing agent {agent.session_id} due to blocked ip: {ip}')
                with Session as db:
                    self.agent_task_service.create_task_exit(db, agent, current_user_id=1)

    def shutdown(self):
        """
        Kills additional processes that were spawned
        """
        # If the plugin spawns a process provide a shutdown method for when Empire exits else leave it as pass
        pass