Files
sbp/tests/test_install.py
T

159 lines
6.6 KiB
Python

import os
import shutil
import tempfile
import unittest
from pathlib import Path
from unittest import mock
from sbp import net, paths as path_module, tls
from sbp.errors import SbpError
from sbp.ops import install as install_op
from sbp.ops import protocols as protocols_op
from sbp.ops import upgrade as upgrade_op
from sbp.ops import users as users_op
from sbp.state import RealityState, StateStore
class InstallTest(unittest.TestCase):
def setUp(self) -> None:
self.tmp = Path(tempfile.mkdtemp(prefix="sbp-test-"))
self.cert = self.tmp / "cert.pem"
self.key = self.tmp / "key.pem"
tls.generate("www.bing.com", self.cert, self.key)
self.env = mock.patch.dict(
os.environ,
{
"SBP_CONFIG_DIR": str(self.tmp / "config"),
"SBP_DATA_DIR": str(self.tmp / "data"),
"SBP_UNIT_DIR": str(self.tmp / "units"),
},
)
self.env.start()
self.patchers = [
mock.patch("sbp.ops.install.hostinfo.check_platform"),
mock.patch("sbp.ops.install.hostinfo.warnings", return_value=[]),
mock.patch("sbp.ops.install.hostinfo.detect_host", return_value="203.0.113.10"),
mock.patch("sbp.ops.install.hostinfo.ipv6_available", return_value=True),
mock.patch("sbp.ops.install.hostinfo.default_listen", return_value="::"),
mock.patch("sbp.ops.install.net.get_release", return_value=net.Release("v1.14.4", [])),
mock.patch("sbp.singbox.install_binary", side_effect=self._install_binary),
mock.patch("sbp.singbox.check_config"),
mock.patch(
"sbp.singbox.reality_keypair",
return_value=RealityState(private_key="PRIVATE", public_key="PUBLIC"),
),
mock.patch("sbp.ops.common.restart"),
mock.patch("sbp.service.enable"),
mock.patch("sbp.service.active", return_value=True),
]
for patcher in self.patchers:
patcher.start()
self.paths = path_module.load()
def tearDown(self) -> None:
for patcher in self.patchers:
patcher.stop()
self.env.stop()
shutil.rmtree(self.tmp, ignore_errors=True)
def _install_binary(self, paths, release, current_config=None, progress=None):
paths.binary.parent.mkdir(parents=True, exist_ok=True)
paths.binary.write_text("#!/bin/sh\nexit 0\n")
os.chmod(paths.binary, 0o755)
def test_install_creates_state_config_and_unit(self) -> None:
state = install_op.install(
self.paths,
install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key), sni="www.bing.com"),
)
self.assertTrue(self.paths.state_file.is_file())
self.assertTrue(self.paths.server_config.is_file())
self.assertTrue(self.paths.unit_file.is_file())
self.assertEqual(len(state.users), 1)
self.assertEqual(set(state.protocols), {"hysteria2", "vless-reality", "trojan", "shadowsocks", "tuic", "anytls"})
ports = [proto.port for proto in state.protocols.values()]
self.assertEqual(len(ports), len(set(ports)))
self.assertTrue(all(port >= 20000 for port in ports))
def test_add_and_remove_user(self) -> None:
install_op.install(self.paths, install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key)))
users_op.add(self.paths, "phone")
store = path_module.load()
state = StateStore(store.config_dir).load()
self.assertEqual([user.name for user in state.users], ["u1", "phone"])
users_op.remove(self.paths, "phone")
state = StateStore(store.config_dir).load()
self.assertEqual([user.name for user in state.users], ["u1"])
def test_cannot_remove_last_user(self) -> None:
install_op.install(self.paths, install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key)))
with self.assertRaises(SbpError):
users_op.remove(self.paths, "u1")
def test_disable_protocol_updates_config(self) -> None:
import json
install_op.install(self.paths, install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key)))
protocols_op.disable(self.paths, "trojan")
config = json.loads(self.paths.server_config.read_text())
tags = {inbound["tag"] for inbound in config["inbounds"]}
self.assertNotIn("trojan-in", tags)
protocols_op.enable(self.paths, "trojan")
config = json.loads(self.paths.server_config.read_text())
tags = {inbound["tag"] for inbound in config["inbounds"]}
self.assertIn("trojan-in", tags)
def test_set_port_rejects_port_of_another_protocol(self) -> None:
state = install_op.install(
self.paths, install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key))
)
taken = state.protocols["trojan"].port
with self.assertRaises(SbpError):
protocols_op.set_port(self.paths, "vless-reality", taken)
def test_upgrade_restores_missing_binary(self) -> None:
install_op.install(
self.paths, install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key))
)
self.paths.binary.unlink()
current, target = upgrade_op.upgrade(self.paths)
self.assertTrue(self.paths.binary.is_file())
self.assertEqual(target, "v1.14.4")
self.assertNotEqual(current, "missing")
def test_set_port_to_current_port_is_noop(self) -> None:
state = install_op.install(
self.paths, install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key))
)
port = state.protocols["trojan"].port
protocols_op.set_port(self.paths, "trojan", port)
reloaded = StateStore(self.paths.config_dir).load()
self.assertEqual(reloaded.protocols["trojan"].port, port)
def test_reinstall_rejects_without_force(self) -> None:
install_op.install(self.paths, install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key)))
with self.assertRaises(Exception):
install_op.install(self.paths, install_op.InstallOptions(host="203.0.113.10"))
def test_reinstall_keeps_custom_certificate(self) -> None:
install_op.install(
self.paths,
install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key), sni="www.bing.com"),
)
state = install_op.install(self.paths, install_op.InstallOptions(force=True))
self.assertEqual(state.tls.mode, "custom")
self.assertEqual(state.tls.cert, str(self.cert.resolve()))
def test_reinstall_keeps_ports(self) -> None:
first = install_op.install(
self.paths,
install_op.InstallOptions(host="203.0.113.10", cert=str(self.cert), key=str(self.key)),
)
ports = {key: proto.port for key, proto in first.protocols.items()}
second = install_op.install(self.paths, install_op.InstallOptions(force=True))
self.assertEqual({key: proto.port for key, proto in second.protocols.items()}, ports)
if __name__ == "__main__":
unittest.main()