warpgate/tests/test_http_websocket.py

82 lines
2.4 KiB
Python
Raw Normal View History

2022-08-14 18:36:49 +08:00
import ssl
import requests
from websocket import create_connection
2022-11-12 00:00:12 +08:00
from uuid import uuid4
2022-08-14 18:36:49 +08:00
2022-11-12 00:00:12 +08:00
from .api_client import (
api_admin_session,
api_create_target,
api_create_user,
api_create_role,
api_add_role_to_user,
api_add_role_to_target,
)
from .conftest import WarpgateProcess
from .test_http_common import * # noqa
2022-08-14 18:36:49 +08:00
class TestHTTPWebsocket:
def test_basic(
self,
2022-11-12 00:00:12 +08:00
echo_server_port,
shared_wg: WarpgateProcess,
2022-08-14 18:36:49 +08:00
):
2022-11-12 00:00:12 +08:00
url = f"https://localhost:{shared_wg.http_port}"
with api_admin_session(url) as session:
role = api_create_role(url, session, {"name": f"role-{uuid4()}"})
user = api_create_user(
url,
session,
{
"username": f"user-{uuid4()}",
"credentials": [
{
"kind": "Password",
"hash": "123",
},
],
},
)
api_add_role_to_user(url, session, user["id"], role["id"])
echo_target = api_create_target(
url,
session,
{
"name": f"echo-{uuid4()}",
"options": {
"kind": "Http",
"url": f"http://localhost:{echo_server_port}",
"tls": {
"mode": "Disabled",
"verify": False,
},
},
},
)
api_add_role_to_target(url, session, echo_target["id"], role["id"])
2022-08-14 18:36:49 +08:00
session = requests.Session()
session.verify = False
session.post(
2022-11-12 00:00:12 +08:00
f"{url}/@warpgate/api/auth/login",
2022-08-14 18:36:49 +08:00
json={
2022-11-12 00:00:12 +08:00
"username": user["username"],
"password": "123",
2022-08-14 18:36:49 +08:00
},
)
cookies = session.cookies.get_dict()
2022-11-12 00:00:12 +08:00
cookie = "; ".join([f"{k}={v}" for k, v in cookies.items()])
2022-08-14 18:36:49 +08:00
ws = create_connection(
2022-11-12 00:00:12 +08:00
f"wss://localhost:{shared_wg.http_port}/socket?warpgate-target={echo_target['name']}",
2022-08-14 18:36:49 +08:00
cookie=cookie,
sslopt={"cert_reqs": ssl.CERT_NONE},
)
2022-11-12 00:00:12 +08:00
ws.send("test")
assert ws.recv() == "test"
ws.send_binary(b"test")
assert ws.recv() == b"test"
2022-08-14 18:36:49 +08:00
ws.ping()
ws.close()