109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
import json
|
|
import threading
|
|
import unittest
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
from app.routeros import RouterOSClient
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
last_put = None
|
|
|
|
def do_GET(self):
|
|
data = b"[]"
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
|
|
def do_PUT(self):
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
Handler.last_put = json.loads(self.rfile.read(length))
|
|
data = json.dumps(Handler.last_put).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
|
|
def log_message(self, fmt, *args):
|
|
return
|
|
|
|
|
|
class UnblockHandler(BaseHTTPRequestHandler):
|
|
deleted_path = None
|
|
|
|
def do_GET(self):
|
|
data = json.dumps([{
|
|
".id": "*1",
|
|
"list": "IDS-BLOCK",
|
|
"address": "9.9.9.9",
|
|
"timeout": "1h",
|
|
"comment": "test",
|
|
}]).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
|
|
def do_DELETE(self):
|
|
UnblockHandler.deleted_path = self.path
|
|
self.send_response(204)
|
|
self.send_header("Content-Length", "0")
|
|
self.end_headers()
|
|
|
|
def log_message(self, fmt, *args):
|
|
return
|
|
|
|
|
|
class RouterOSTests(unittest.TestCase):
|
|
def test_put_address_list_entry(self):
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
client = RouterOSClient(
|
|
f"http://127.0.0.1:{server.server_port}",
|
|
"user",
|
|
"secret",
|
|
False,
|
|
"IDS-BLOCK",
|
|
2,
|
|
)
|
|
result = client.block_ip("9.9.9.9", "1h", "test")
|
|
self.assertTrue(result.success)
|
|
self.assertEqual(Handler.last_put["list"], "IDS-BLOCK")
|
|
self.assertEqual(Handler.last_put["address"], "9.9.9.9")
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
|
|
def test_list_and_unblock_address_list_entry(self):
|
|
UnblockHandler.deleted_path = None
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0), UnblockHandler)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
client = RouterOSClient(
|
|
f"http://127.0.0.1:{server.server_port}",
|
|
"user",
|
|
"secret",
|
|
False,
|
|
"IDS-BLOCK",
|
|
2,
|
|
)
|
|
rows = client.list_blocks()
|
|
self.assertEqual(rows[0]["address"], "9.9.9.9")
|
|
result = client.unblock_ip("9.9.9.9")
|
|
self.assertTrue(result.success)
|
|
self.assertTrue(UnblockHandler.deleted_path.endswith("/*1"))
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|