更新

添加请求全量封禁IP方法

Author: SecAegis

device/sdk/python/README.md

@@ -61,14 +61,55 @@ sec_auto_ban.run()
 
 ## 参数说明
 
-| 参数           | 描述                    | 是否需要填写          |
-| ---------------- | ------------------------- | --------------- |
-| server_ip        | 核心模块回连IP      | 需要             |
-| server_port      | 核心模块回连端口  | 需要             |
-| sk               | 设备页面生成的密钥 | 需要             |
-| client_type      | 模块类型(`alarm`/`block`) | 需要             |
-| enable_cidr      | 封禁模块是否开启 `Cidr` 封禁，若开启`block_ip()`和`unblock_ip()`参数将传入`Cidr` | 可选，默认为 `False`             |
-| alarm_analysis   | 告警分析函数        | `alarm`模块必填 |
-| block_ip         | 封禁函数              | `block`模块必填 |
-| unblock_ip       | 解禁函数              | `block`模块必填 |
-| get_all_block_ip | 获取设备中全部封禁IP函数 | `block`模块可选 |
+| 参数                     | 描述                                                           | 是否需要填写         |
+|------------------------|--------------------------------------------------------------|----------------|
+| server_ip              | 核心模块回连IP                                                     | 需要             |
+| server_port            | 核心模块回连端口                                                     | 需要             |
+| sk                     | 设备页面生成的密钥                                                    | 需要             |
+| client_type            | 模块类型(`alarm`/`block`)                                        | 需要             |
+| enable_cidr            | 封禁模块是否开启 `Cidr` 封禁，若开启`block_ip()`和`unblock_ip()`参数将传入`Cidr` | 可选，默认为 `False` |
+| alarm_analysis         | 告警分析函数                                                       | `alarm`模块必填    |
+| block_ip               | 封禁函数                                                         | `block`模块必填    |
+| unblock_ip             | 解禁函数                                                         | `block`模块必填    |
+| get_all_block_ip       | 获取设备中全部封禁IP函数                                                | `block`模块可选    |
+| login_success_callback | 登陆成功回调                                                       | 可选             |
+
+## SDK调用方法
+
+### send_alarm()
+
+告警设备向平台发送告警信息。
+
+eg:
+
+```python
+def alarm_analysis(ws_client):
+    ws_client.send_alarm("攻击IP", "被攻击资产", "攻击方式")
+```
+
+or:
+
+```
+sec_auto_ban.send_alarm("攻击IP", "被攻击资产", "攻击方式")
+```
+
+### send_notify()
+
+向平台发送通知。
+
+eg:
+
+```python
+sec_auto_ban.send_notify("封禁失败", "xxx设备无法连接服务器")
+```
+
+### send_sync()
+
+封禁设备主动向平台请求全部封禁IP。常用于脚本第一次启动，需同步全量IP场景。
+
+eg:
+
+```python
+def login_success_callback():
+    sec_auto_ban.send_sync()
+```

device/sdk/python/SecAutoBan/sec_auto_ban.py

@@ -5,19 +5,28 @@
 class SecAutoBan:
     init = False
 
-    def __init__(self, server_ip, server_port, sk, client_type, enable_cidr=False, alarm_analysis=None, block_ip=None, unblock_ip=None, get_all_block_ip=None):
+    def __init__(self, server_ip, server_port, sk, client_type, enable_cidr=False, alarm_analysis=None, block_ip=None, unblock_ip=None, get_all_block_ip=None, login_success_callback=None):
         self.client_type = client_type
         if client_type == "alarm":
             if alarm_analysis is None:
                 util.print("[-] 初始化失败: 未实现告警函数")
                 return
             self.alarm_analysis = alarm_analysis
-        self.ws_client = WebSocketClient(server_ip, server_port, sk, client_type, enable_cidr, block_ip, unblock_ip, get_all_block_ip)
+        self.ws_client = WebSocketClient(server_ip, server_port, sk, client_type, enable_cidr, block_ip, unblock_ip, get_all_block_ip, login_success_callback)
         self.init = True
 
     def print(self, message):
         util.print(message)
 
+    def send_alarm(self, ip: str, attack_asset: str, attack_method: str):
+        self.ws_client.send_alarm(ip, attack_asset, attack_method)
+
+    def send_notify(self, title: str, content: str):
+        self.ws_client.send_notify(title, content)
+
+    def send_sync(self):
+        self.ws_client.send_sync()
+
     def run(self):
         if not self.init:
             return

device/sdk/python/SecAutoBan/websocket_client.py

@@ -13,7 +13,7 @@ class WebSocketClient:
     send_alarm_ip_list = []
     pool = ThreadPool(processes=2)
 
-    def __init__(self, server_ip: str, server_port: int, sk: str, client_type: str, enable_cidr=False, block_ip=None, unblock_ip=None, get_all_block_ip=None):
+    def __init__(self, server_ip: str, server_port: int, sk: str, client_type: str, enable_cidr=False, block_ip=None, unblock_ip=None, get_all_block_ip=None, login_success_callback=None):
         self.server_ip = server_ip
         self.server_port = server_port
         self.sk = sk
@@ -29,6 +29,7 @@ def __init__(self, server_ip: str, server_port: int, sk: str, client_type: str,
             self.unblock_ip = unblock_ip
             self.get_all_block_ip = get_all_block_ip
         self.client_type = client_type
+        self.login_success_callback = login_success_callback
         self.ws = websocket.WebSocketApp(
             "ws://" + server_ip + ":" + str(server_port) + "/device",
             on_message=self.on_message,
@@ -60,6 +61,8 @@ def on_message(self, w, message):
         if message["method"] == "login":
             self.is_login = True
             util.print("[+] 登录成功，设备名称: " + message["data"]["deviceName"])
+            if self.login_success_callback is not None:
+                self.login_success_callback()
         if self.client_type == "block":
             if message["method"] == "blockCidr":
                 util.print("[+] 封禁IP: " + message["data"]["cidr"])
@@ -124,9 +127,9 @@ def connect(self):
             return
         self.pool.apply_async(self.web_socket_d)
 
-    def send_alarm(self, ip: str, attackAsset: str, attackMethod: str):
-        if self.client_type == "block":
-            util.print("[-] 封禁模块无法发送告警数据")
+    def send_alarm(self, ip: str, attack_asset: str, attack_method: str):
+        if self.client_type != "alarm":
+            util.print("[-] 非告警模块无法发送告警数据")
             return
         if not self.is_login:
             util.print("[-] 未登录成功，无法发送数据")
@@ -140,12 +143,12 @@ def send_alarm(self, ip: str, attackAsset: str, attackMethod: str):
             "method": "alarmIp",
             "data": {
                 "ip": ip,
-                "attackAsset": attackAsset,
-                "attackMethod": attackMethod
+                "attackAsset": attack_asset,
+                "attackMethod": attack_method
             }
         }
         iv = util.random_bytes()
-        util.print("[+] 发送告警IP: " + ip + "->" + attackAsset + "\t" + attackMethod)
+        util.print("[+] 发送告警IP: " + ip + "->" + attack_asset + "\t" + attack_method)
         self.ws.send(iv + util.aes_cfb_encrypt(self.sk[3:].encode(), iv, json.dumps(send_data).encode()))
     def send_notify(self, title: str, content: str):
         if not self.is_login:
@@ -159,4 +162,16 @@ def send_notify(self, title: str, content: str):
             }
         }
         iv = util.random_bytes()
+        self.ws.send(iv + util.aes_cfb_encrypt(self.sk[3:].encode(), iv, json.dumps(send_data).encode()))
+    def send_sync(self):
+        if self.client_type != "block":
+            util.print("[-] 非封禁模块无法请求封禁IP")
+            return
+        if not self.is_login:
+            util.print("[-] 未登录成功，无法发送数据")
+            return
+        send_data = {
+            "method": "syncBlockIp"
+        }
+        iv = util.random_bytes()
         self.ws.send(iv + util.aes_cfb_encrypt(self.sk[3:].encode(), iv, json.dumps(send_data).encode()))

device/sdk/python/setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="SecAutoBan",
-    version="4.0.3",
+    version="5.0.0",
     author="SecReport",
     author_email="secaegis@outlook.com",
     description="SecAutoBan SDK",