✨ add module write database

Author: TuuuNya

WebPocket.py

@@ -1,4 +1,4 @@
-from lib.pocket import Pocket
+from lib.Pocket import Pocket
 
 
 class WebPocket(Pocket):

database/pocket.db

@@ -17,14 +17,10 @@ def __init__(self):
         pass
 
     def get_info(self):
-        return {
-            "name": self.name,
-            "description": self.description,
-            "author": self.author,
-            "references": self.references,
-            "service_name": self.service_name,
-            "service_version": self.service_version,
-        }
+        info = {}
+        for field_name in self.info_fields:
+            info[field_name] = getattr(self, field_name)
+        return info
 
     def update_info(self, info):
         for name in info:

lib/BaseExploit.py

@@ -0,0 +1,79 @@
+import os
+import sqlite3
+from fnmatch import fnmatchcase
+from utils.files import ROOT_PATH
+from utils.module import name_convert
+from importlib import import_module
+
+
+class Database:
+    db_file = '{root_path}/database/pocket.db'.format(root_path=ROOT_PATH)
+    connection = None
+    cursor = None
+
+    def __init__(self):
+        self.connection = sqlite3.connect(self.db_file)
+        self.cursor = self.connection.cursor()
+
+        self.create_table()
+
+        # 初始化数据
+        if self.get_module_count() == 0:
+            self.db_rebuild()
+
+    def get_module_count(self):
+        sql = 'select count(*) from modules;'
+        rs = self.cursor.execute(sql)
+        (count, ) = rs.fetchone()
+        return count
+
+    def create_table(self):
+        init_table_sql = (
+            'CREATE TABLE IF NOT EXISTS "modules" ('
+            '"id" INTEGER NOT NULL,'
+            '"name" TEXT,'
+            '"module_name" TEXT,'
+            '"description" TEXT,'
+            '"author" TEXT,'
+            '"references" TEXT,'
+            '"disclosure_date" TEXT,'
+            '"service_name" TEXT,'
+            '"service_version" TEXT,'
+            'PRIMARY KEY("id")'
+            ');'
+        )
+        self.cursor.execute(init_table_sql)
+
+    def delete_table(self):
+        delete_table_sql = "delete from modules;"
+        self.cursor.execute(delete_table_sql)
+
+    def insert_module(self, info):
+        with self.connection:
+            self.connection.execute(
+                "insert into modules \
+                (name, module_name, description, author, 'references', disclosure_date, service_name, service_version) \
+                values (?, ?, ?, ?, ?, ?, ?, ?)",
+                (info.get('name'), info.get('module_name'), info.get('description'), '|'.join(info.get('author')),
+                 '|'.join(info.get('references')), info.get('disclosure_date'), info.get('service_name'),
+                 info.get('service_version'))
+            )
+
+    def db_rebuild(self):
+        self.delete_table()
+        self.create_table()
+
+        for directory_name, directories, filenames in os.walk('modules/'):
+            for filename in filenames:
+                if filename not in ['__init__.py']\
+                        and not fnmatchcase(filename, "*.pyc")\
+                        and fnmatchcase(filename, "*.py"):
+                    full_name = "{directory}/{filename}".format(directory=directory_name, filename=filename)
+                    module_name = name_convert(full_name)
+                    module_class = import_module("modules.{module_name}".format(
+                        module_name=module_name.replace("/", ".")
+                    ))
+                    module_instance = module_class.Exploit()
+                    module_info = module_instance.get_info()
+                    module_info['module_name'] = module_name
+                    self.insert_module(module_info)

lib/Database.py

@@ -1,14 +1,16 @@
 from lib.cmd2 import Cmd
 from utils.files import ROOT_PATH
+from utils.module import name_convert
 from pathlib import Path
 from colorama import Fore, Style
 from tabulate import tabulate
 from importlib import import_module
+from lib.Database import Database
 from lib.ExploitOption import ExploitOption
 from lib.exception.Module import ModuleNotUseException
 
 
-class Pocket(Cmd):
+class Pocket(Cmd, Database):
     colors = "Always"
 
     console_prompt = "{COLOR_START}WebPocket{COLOR_END}".format(COLOR_START="\033[4m", COLOR_END="\033[0m")
@@ -17,7 +19,8 @@ class Pocket(Cmd):
     module_instance = None
 
     def __init__(self):
-        super().__init__()
+        super(Pocket, self).__init__()
+        Database.__init__(self)
         self.prompt = self.console_prompt + self.console_prompt_end
         self.hidden_commands.extend(['alias', 'edit', 'macro', 'py', 'pyscript', 'shell', 'shortcuts', 'load'])
 
@@ -32,7 +35,7 @@ def do_set(self, args):
         self.module_instance.options.set_option(arg, value)
 
     def do_use(self, module_name):
-        module_file = "{ROOT}/modules/{MODULE}.py".format(ROOT=ROOT_PATH, MODULE=module_name)
+        module_file = name_convert(module_name)
         module_type = module_name.split("/")[0]
 
         if Path(module_file).is_file():
@@ -94,6 +97,10 @@ def do_exploit(self, args):
             style_end=Style.RESET_ALL
         ))
 
+    def do_db_rebuild(self, args):
+        self.db_rebuild()
+        self.poutput("Database rebuild done.", color=Fore.GREEN)
+
     def set_prompt(self, module_type, module_name):
         module_prompt = " {module_type}({color}{module_name}{color_end})".format(
             module_type=module_type,

lib/pocket.py lib/Pocket.py

@@ -0,0 +1,36 @@
+from lib.BaseExploit import BaseExploit
+from lib.ExploitOption import ExploitOption
+
+
+class Exploit(BaseExploit):
+
+    def __init__(self):
+        super(Exploit, self).__init__()
+        self.update_info(info={
+            "name": "zabbix 2.0.3 sqli",
+            "description": "zabbix 2.0.3 jsrpc.php sqli",
+            "author": ["TuuuNya", "unknown"],
+            "references": [
+                "https://www.hackersb.cn/hacker/219.html",
+                "https://www.hackersb.cn/hacker/167.html",
+            ],
+            "disclosure_date": "2016-08-22",
+            "service_name": "zabbix",
+            "service_version": "2.0.6",
+        })
+        self.register_options([
+            ExploitOption(
+                name="host",
+                required=True,
+                description="The target domain",
+                value=None
+            ),
+        ])
+
+    def check(self):
+        pass
+
+    def exploit(self):
+        super(Exploit, self).exploit()
+        print(self.options.get_option("host"))
+        return "Exploit success, webshell: http://www.hackersb.cn/a.php"

modules/exploits/cms/zabbix_2_0_3_sqli.py

@@ -0,0 +1,10 @@
+from utils.files import ROOT_PATH
+
+
+def name_convert(name):
+    if name.find(".py") is not -1:
+        module_name = name.replace("modules/", "").replace(".py", "")
+        return module_name
+    else:
+        full_name = "{ROOT}/modules/{MODULE}.py".format(ROOT=ROOT_PATH, MODULE=name)
+        return full_name