From 7900e8b579b8a0f92fa03736f28ebcd2bce004f9 Mon Sep 17 00:00:00 2001 From: pohmelie Date: Wed, 3 Jan 2024 23:07:38 +0400 Subject: [PATCH] refactor project structure --- .coveragerc | 4 -- .flake8 | 2 - .github/workflows/ci.yml | 48 ++++++++--------- .pre-commit-config.yaml | 50 +++++++++++++++++ MANIFEST.in | 4 -- examples/shadowsocks-like.py | 20 ++++--- pyproject.toml | 66 +++++++++++++++++++++++ pytest.ini | 7 --- readme.md | 2 +- setup.cfg | 28 ---------- setup.py | 4 -- siosocks/__init__.py | 2 - src/siosocks/__init__.py | 4 ++ {siosocks => src/siosocks}/__main__.py | 49 +++++++++++------ {siosocks => src/siosocks}/exceptions.py | 0 {siosocks => src/siosocks}/interface.py | 1 - {siosocks => src/siosocks}/io/__init__.py | 0 {siosocks => src/siosocks}/io/asyncio.py | 40 ++++++++++---- {siosocks => src/siosocks}/io/const.py | 0 {siosocks => src/siosocks}/io/socket.py | 9 ++-- {siosocks => src/siosocks}/io/trio.py | 40 ++++++++++---- {siosocks => src/siosocks}/protocol.py | 36 ++++++------- {siosocks => src/siosocks}/sansio.py | 4 +- tests/test_asyncio.py | 29 +++++++--- tests/test_protocol.py | 51 +++--------------- tests/test_sansio.py | 5 +- tests/test_socketserver.py | 22 +++++--- tests/test_trio.py | 23 +++++--- 28 files changed, 325 insertions(+), 225 deletions(-) delete mode 100644 .coveragerc delete mode 100644 .flake8 create mode 100644 .pre-commit-config.yaml delete mode 100644 MANIFEST.in create mode 100644 pyproject.toml delete mode 100644 pytest.ini delete mode 100644 setup.cfg delete mode 100644 setup.py delete mode 100644 siosocks/__init__.py create mode 100644 src/siosocks/__init__.py rename {siosocks => src/siosocks}/__main__.py (83%) rename {siosocks => src/siosocks}/exceptions.py (100%) rename {siosocks => src/siosocks}/interface.py (99%) rename {siosocks => src/siosocks}/io/__init__.py (100%) rename {siosocks => src/siosocks}/io/asyncio.py (78%) rename {siosocks => src/siosocks}/io/const.py (100%) rename {siosocks => src/siosocks}/io/socket.py (96%) rename {siosocks => src/siosocks}/io/trio.py (75%) rename {siosocks => src/siosocks}/protocol.py (94%) rename {siosocks => src/siosocks}/sansio.py (99%) diff --git a/.coveragerc b/.coveragerc deleted file mode 100644 index 3282d6a..0000000 --- a/.coveragerc +++ /dev/null @@ -1,4 +0,0 @@ -[report] -show_missing = True -omit = - siosocks/__main__.py diff --git a/.flake8 b/.flake8 deleted file mode 100644 index 6deafc2..0000000 --- a/.flake8 +++ /dev/null @@ -1,2 +0,0 @@ -[flake8] -max-line-length = 120 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b31b257..58607b9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,40 +6,40 @@ jobs: lint: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 - with: - python-version: "3.x" - - run: pip install flake8 - - run: flake8 + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + with: + python-version: "3.11" + - run: pip install -e ./[dev] + - run: pre-commit run -a tests: needs: lint runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.8", "3.9", "3.10"] + python-version: ["3.11", "3.12"] steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - - run: pip install -e ./[dev] - - run: pytest - - uses: codecov/codecov-action@v2 - with: - fail_ci_if_error: true - verbose: true + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + - run: pip install -e ./[dev] + - run: pytest + - uses: codecov/codecov-action@v2 + with: + fail_ci_if_error: true + verbose: true deploy: needs: tests runs-on: ubuntu-latest if: github.ref == 'refs/heads/master' steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 - - uses: casperdcl/deploy-pypi@v2 - with: - password: ${{ secrets.PYPI_TOKEN }} - build: true - skip_existing: true + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + - uses: casperdcl/deploy-pypi@v2 + with: + password: ${{ secrets.PYPI_TOKEN }} + build: true + skip_existing: true diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..e93248b --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,50 @@ +# See https://pre-commit.com for more information +# See https://pre-commit.com/hooks.html for more hooks + +.python-linters: &python-linters + pass_filenames: false + fail_fast: true + language: system + types: [python] + +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.5.0 + hooks: + - id: check-ast + fail_fast: true + - id: trailing-whitespace + - id: check-toml + fail_fast: true + - id: end-of-file-fixer + fail_fast: true + +- repo: https://github.com/asottile/add-trailing-comma + rev: v2.1.0 + hooks: + - id: add-trailing-comma + fail_fast: true + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.1.0 + hooks: + - id: pretty-format-yaml + fail_fast: true + args: + - --autofix + - --preserve-quotes + - --indent=2 + +- repo: local + hooks: + - <<: *python-linters + id: black + name: Format with Black + entry: black + args: ["."] + + - <<: *python-linters + id: ruff + name: Check with ruff + entry: ruff + args: ["check", "--fix", "."] diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index 1976e63..0000000 --- a/MANIFEST.in +++ /dev/null @@ -1,4 +0,0 @@ -include readme.md -include license.txt -include history.md -recursive-include tests * diff --git a/examples/shadowsocks-like.py b/examples/shadowsocks-like.py index 3306418..224a789 100644 --- a/examples/shadowsocks-like.py +++ b/examples/shadowsocks-like.py @@ -4,9 +4,9 @@ import functools import socket -from siosocks.io.asyncio import ServerIO, ClientIO from siosocks.interface import async_engine -from siosocks.protocol import SocksServer, SocksClient +from siosocks.io.asyncio import ClientIO, ServerIO +from siosocks.protocol import SocksClient, SocksServer # This is very strong encryption method, believe me! @@ -19,7 +19,6 @@ def decode(data: bytes): class CommonProxy: - def __init__(self, proxied): self._proxied = proxied @@ -28,19 +27,16 @@ def __getattr__(self, name): class IncomingDecoder(CommonProxy): - async def read(self, count): return decode(await self._proxied.read(count)) class OutgoingEncoder(CommonProxy): - def write(self, data): return self._proxied.write(encode(data)) class RemoteIO(ServerIO): - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.incoming_reader = IncomingDecoder(self.incoming_reader) @@ -58,17 +54,19 @@ async def open_connection(host=None, port=None, *, socks_host=None, socks_port=N class LocalIO(ServerIO): - def __init__(self, *args, remote_host, remote_port, **kwargs): super().__init__(*args, **kwargs) self.__remote_host = remote_host self.__remote_port = remote_port async def connect(self, host, port): - self.outgoing_reader, self.outgoing_writer = await open_connection(host, port, - socks_host=self.__remote_host, - socks_port=self.__remote_port, - socks_version=5) + self.outgoing_reader, self.outgoing_writer = await open_connection( + host, + port, + socks_host=self.__remote_host, + socks_port=self.__remote_port, + socks_version=5, + ) async def socks_server_handler(reader, writer, *, io_factory, **kwargs): diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..29a2746 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,66 @@ +[project] +name = "siosocks" +version = "0.3.0" +description = "sans-io socks proxy client/server with couple io backends" +readme = "readme.md" +requires-python = ">= 3.11" +license = {file = "license.txt"} +authors = [ + {name = "pohmelie", email = "multisosnooley@gmail.com"}, +] +classifiers = [ + "Programming Language :: Python", + "Programming Language :: Python :: 3", +] + +[project.urls] +Github = "https://github.com/pohmelie/siosocks" + +[project.optional-dependencies] +dev = [ + # tests + "pytest-asyncio", + "pytest-cov", + "pytest-trio", + "pytest", + "trio", + + # linters + "pre-commit", + "black", + "ruff", +] +trio = [ + "trio", +] + +[build-system] +requires = ["setuptools", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages.find.where = ["src"] + +# tools +[tool.black] +line-length = 120 +target-version = ["py311"] + +[tool.ruff] +line-length = 120 +target-version = "py311" +select = ["E", "W", "F", "Q", "UP", "I", "ASYNC"] +src = ["src"] + +[tool.coverage] +run.source = ["./src/siosocks"] +run.omit = ["./src/siosocks/__main__.py"] +report.show_missing = true + +[tool.pytest.ini_options] +addopts = "-x --durations 10 -p no:anyio --cov" +testpaths = "tests" +log_format = "%(asctime)s.%(msecs)03d %(name)-20s %(levelname)-8s %(filename)-15s %(lineno)-4d %(message)s" +log_date_format = "%H:%M:%S" +log_level = "DEBUG" +asyncio_mode = "strict" diff --git a/pytest.ini b/pytest.ini deleted file mode 100644 index f935feb..0000000 --- a/pytest.ini +++ /dev/null @@ -1,7 +0,0 @@ -[pytest] -addopts = -x --durations=10 --cov-config=.coveragerc --cov=siosocks --cov-report=xml --cov-report=term --cov-report=term-missing -testpaths = tests -log_format = %(asctime)s.%(msecs)03d %(name)-20s %(levelname)-8s %(filename)-15s %(lineno)-4d %(message)s -log_date_format = %H:%M:%S -log_level = DEBUG -asyncio_mode = strict diff --git a/readme.md b/readme.md index 7444c94..9a0717c 100644 --- a/readme.md +++ b/readme.md @@ -14,7 +14,7 @@ - Fun # Features -- Only tcp connect (no bind, no udp) +- Only tcp connect (no bind or udp associate) - Both client and server - Socks versions: 4, 4a, 5 - Socks5 auth: no auth, username/password diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index 7f45270..0000000 --- a/setup.cfg +++ /dev/null @@ -1,28 +0,0 @@ -[metadata] -name = siosocks -version = attr: siosocks.__version__ -url = https://github.com/pohmelie/siosocks -author = pohmelie -author_email = multisosnooley@gmail.com -description = sans-io socks proxy client/server with couple io backends -long_description = file: readme.md -long_description_content_type = text/markdown -license = MIT -license_file = license.txt -classifiers = - Programming Language :: Python - Programming Language :: Python :: 3 - -[options] -packages = find: -python_requires = >= 3.8 - -[options.extras_require] -dev = - pytest - pytest-cov - pytest-asyncio - trio - pytest-trio -trio = - trio diff --git a/setup.py b/setup.py deleted file mode 100644 index b024da8..0000000 --- a/setup.py +++ /dev/null @@ -1,4 +0,0 @@ -from setuptools import setup - - -setup() diff --git a/siosocks/__init__.py b/siosocks/__init__.py deleted file mode 100644 index a35c890..0000000 --- a/siosocks/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -__version__ = "0.3.0" -version = tuple(map(int, __version__.split("."))) diff --git a/src/siosocks/__init__.py b/src/siosocks/__init__.py new file mode 100644 index 0000000..724dda2 --- /dev/null +++ b/src/siosocks/__init__.py @@ -0,0 +1,4 @@ +import importlib.metadata + +__version__ = importlib.metadata.version(__package__) +version = tuple(map(int, __version__.split("."))) diff --git a/siosocks/__main__.py b/src/siosocks/__main__.py similarity index 83% rename from siosocks/__main__.py rename to src/siosocks/__main__.py index ddc1239..a0f095c 100644 --- a/siosocks/__main__.py +++ b/src/siosocks/__main__.py @@ -1,32 +1,47 @@ import argparse -import sys +import asyncio +import contextlib import functools import socket -import asyncio import socketserver -import contextlib +import sys from . import __version__ from .io.asyncio import socks_server_handler as asyncio_socks_server_handler from .io.socket import socks_server_handler as socket_socks_server_handler from .protocol import DEFAULT_ENCODING - parser = argparse.ArgumentParser("siosocks", description="Socks proxy server") -parser.add_argument("--backend", default="asyncio", choices=["asyncio", "socketserver", "trio"], - help="Socks server backend [default: %(default)s]") +parser.add_argument( + "--backend", + default="asyncio", + choices=["asyncio", "socketserver", "trio"], + help="Socks server backend [default: %(default)s]", +) parser.add_argument("--host", default=None, help="Socks server host [default: %(default)s]") parser.add_argument("--port", default=1080, type=int, help="Socks server port [default: %(default)s]") -parser.add_argument("--family", choices=("ipv4", "ipv6", "auto"), default="auto", - help="Socket family [default: %(default)s]") -parser.add_argument("--socks", action="append", type=int, default=[], - help="Socks protocol version [default: %(default)s]") +parser.add_argument( + "--family", + choices=("ipv4", "ipv6", "auto"), + default="auto", + help="Socket family [default: %(default)s]", +) +parser.add_argument( + "--socks", + action="append", + type=int, + default=[], + help="Socks protocol version [default: %(default)s]", +) parser.add_argument("--username", default=None, help="Socks auth username [default: %(default)s]") parser.add_argument("--password", default=None, help="Socks auth password [default: %(default)s]") parser.add_argument("--encoding", default=DEFAULT_ENCODING, help="String encoding [default: %(default)s]") -parser.add_argument("--no-strict", default=False, action="store_true", - help="Allow multiversion socks server, when socks5 used with username/password auth " - "[default: %(default)s]") +parser.add_argument( + "--no-strict", + default=False, + action="store_true", + help="Allow multiversion socks server, when socks5 used with username/password auth " "[default: %(default)s]", +) parser.add_argument("-v", "--version", action="store_true", help="Show siosocks version") ns = parser.parse_args() if ns.version: @@ -40,13 +55,14 @@ }[ns.family] socks_versions = set(ns.socks) or {4, 5} if 4 in socks_versions and ns.username is not None: - print("Socks4 do not provide auth methods, but socks4 allowed " - "and auth required and strict security policy enabled") + print( + "Socks4 do not provide auth methods, but socks4 allowed " + "and auth required and strict security policy enabled", + ) sys.exit(1) def asyncio_main(socks_versions, family, ns): - async def main(): server = await asyncio.start_server(handler, host=ns.host, port=ns.port, family=family) addresses = [] @@ -91,6 +107,7 @@ def socketserver_main(socks_versions, family, ns): def trio_main(socks_versions, family, ns): import trio + from .io.trio import socks_server_handler as trio_socks_server_handler async def main(): diff --git a/siosocks/exceptions.py b/src/siosocks/exceptions.py similarity index 100% rename from siosocks/exceptions.py rename to src/siosocks/exceptions.py diff --git a/siosocks/interface.py b/src/siosocks/interface.py similarity index 99% rename from siosocks/interface.py rename to src/siosocks/interface.py index f6655f2..35ab987 100644 --- a/siosocks/interface.py +++ b/src/siosocks/interface.py @@ -4,7 +4,6 @@ class AbstractSocksIO(abc.ABC): - @abc.abstractmethod def read(self): """ diff --git a/siosocks/io/__init__.py b/src/siosocks/io/__init__.py similarity index 100% rename from siosocks/io/__init__.py rename to src/siosocks/io/__init__.py diff --git a/siosocks/io/asyncio.py b/src/siosocks/io/asyncio.py similarity index 78% rename from siosocks/io/asyncio.py rename to src/siosocks/io/asyncio.py index cf930a3..c2a5eeb 100644 --- a/siosocks/io/asyncio.py +++ b/src/siosocks/io/asyncio.py @@ -3,15 +3,13 @@ from ..exceptions import SocksException from ..interface import AbstractSocksIO, async_engine -from ..protocol import SocksServer, SocksClient, DEFAULT_ENCODING +from ..protocol import DEFAULT_ENCODING, SocksClient, SocksServer from .const import DEFAULT_BLOCK_SIZE - logger = logging.getLogger(__name__) class ServerIO(AbstractSocksIO): - def __init__(self, reader, writer): self.incoming_reader = reader self.incoming_writer = writer @@ -69,7 +67,6 @@ async def socks_server_handler(reader, writer, **kwargs): class ClientIO(AbstractSocksIO): - def __init__(self, reader, writer): self.r = reader self.w = writer @@ -88,19 +85,40 @@ async def passthrough(self): return -async def open_connection(host=None, port=None, *, socks_host=None, socks_port=None, - socks_version=None, username=None, password=None, encoding=DEFAULT_ENCODING, - socks4_extras={}, socks5_extras={}, **open_connection_extras): +async def open_connection( + host=None, + port=None, + *, + socks_host=None, + socks_port=None, + socks_version=None, + username=None, + password=None, + encoding=DEFAULT_ENCODING, + socks4_extras={}, + socks5_extras={}, + **open_connection_extras, +): socks_required = socks_host, socks_port, socks_version socks_enabled = all(socks_required) socks_disabled = not any(socks_required) if socks_enabled == socks_disabled: - raise SocksException("Partly passed socks required arguments: " - "socks_host = {!r}, socks_port = {!r}, socks_version = {!r}".format(*socks_required)) + raise SocksException( + "Partly passed socks required arguments: " + "socks_host = {!r}, socks_port = {!r}, socks_version = {!r}".format(*socks_required), + ) if socks_enabled: reader, writer = await asyncio.open_connection(socks_host, socks_port, **open_connection_extras) - protocol = SocksClient(host, port, socks_version, username=username, password=password, encoding=encoding, - socks4_extras=socks4_extras, socks5_extras=socks5_extras) + protocol = SocksClient( + host, + port, + socks_version, + username=username, + password=password, + encoding=encoding, + socks4_extras=socks4_extras, + socks5_extras=socks5_extras, + ) io = ClientIO(reader, writer) await async_engine(protocol, io) else: diff --git a/siosocks/io/const.py b/src/siosocks/io/const.py similarity index 100% rename from siosocks/io/const.py rename to src/siosocks/io/const.py diff --git a/siosocks/io/socket.py b/src/siosocks/io/socket.py similarity index 96% rename from siosocks/io/socket.py rename to src/siosocks/io/socket.py index f0c6bee..00d84cd 100644 --- a/siosocks/io/socket.py +++ b/src/siosocks/io/socket.py @@ -1,13 +1,12 @@ -import socketserver -import socket import logging -from concurrent.futures import ThreadPoolExecutor, wait, FIRST_EXCEPTION +import socket +import socketserver +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait from ..interface import AbstractSocksIO, sync_engine from ..protocol import SocksServer from .const import DEFAULT_BLOCK_SIZE - logger = logging.getLogger(__name__) @@ -15,7 +14,6 @@ class ServerIO(AbstractSocksIO): - def __init__(self, socket): self.incoming_socket = socket self.incoming_socket.settimeout(TIMEOUT) @@ -58,7 +56,6 @@ def _sink(self, producer, consumer): class socks_server_handler(socketserver.BaseRequestHandler): - def __init__(self, *args, socks_protocol_kw, **kwargs): self._socks_protocol_kw = socks_protocol_kw super().__init__(*args, **kwargs) diff --git a/siosocks/io/trio.py b/src/siosocks/io/trio.py similarity index 75% rename from siosocks/io/trio.py rename to src/siosocks/io/trio.py index 4421562..eb943df 100644 --- a/siosocks/io/trio.py +++ b/src/siosocks/io/trio.py @@ -4,15 +4,13 @@ from ..exceptions import SocksException from ..interface import AbstractSocksIO, async_engine -from ..protocol import SocksServer, SocksClient, DEFAULT_ENCODING +from ..protocol import DEFAULT_ENCODING, SocksClient, SocksServer from .const import DEFAULT_BLOCK_SIZE - logger = logging.getLogger(__name__) class ServerIO(AbstractSocksIO): - def __init__(self, stream): self.incoming_stream = stream self.outgoing_stream = None @@ -58,7 +56,6 @@ async def socks_server_handler(stream, **kwargs): class ClientIO(AbstractSocksIO): - def __init__(self, stream): self.stream = stream @@ -76,19 +73,40 @@ async def passthrough(self): return -async def open_tcp_stream(host, port, *, socks_host=None, socks_port=None, - socks_version=None, username=None, password=None, encoding=DEFAULT_ENCODING, - socks4_extras={}, socks5_extras={}, **open_tcp_stream_extras): +async def open_tcp_stream( + host, + port, + *, + socks_host=None, + socks_port=None, + socks_version=None, + username=None, + password=None, + encoding=DEFAULT_ENCODING, + socks4_extras={}, + socks5_extras={}, + **open_tcp_stream_extras, +): socks_required = socks_host, socks_port, socks_version socks_enabled = all(socks_required) socks_disabled = not any(socks_required) if socks_enabled == socks_disabled: - raise SocksException("Partly passed socks required arguments: " - "socks_host = {!r}, socks_port = {!r}, socks_version = {!r}".format(*socks_required)) + raise SocksException( + "Partly passed socks required arguments: " + "socks_host = {!r}, socks_port = {!r}, socks_version = {!r}".format(*socks_required), + ) if socks_enabled: stream = await trio.open_tcp_stream(socks_host, socks_port, **open_tcp_stream_extras) - protocol = SocksClient(host, port, socks_version, username=username, password=password, encoding=encoding, - socks4_extras=socks4_extras, socks5_extras=socks5_extras) + protocol = SocksClient( + host, + port, + socks_version, + username=username, + password=password, + encoding=encoding, + socks4_extras=socks4_extras, + socks5_extras=socks5_extras, + ) io = ClientIO(stream) await async_engine(protocol, io) else: diff --git a/siosocks/protocol.py b/src/siosocks/protocol.py similarity index 94% rename from siosocks/protocol.py rename to src/siosocks/protocol.py index 6e7e750..00f9a33 100644 --- a/siosocks/protocol.py +++ b/src/siosocks/protocol.py @@ -1,12 +1,11 @@ import abc -import enum import contextlib +import enum from ipaddress import IPv4Address, IPv6Address from .exceptions import SocksException from .sansio import SansIORW - DEFAULT_ENCODING = "utf-8" @@ -21,7 +20,6 @@ class SocksCommand(enum.IntEnum): class AbstractSocks(abc.ABC): - def __init__(self, io): self.io = io @@ -44,12 +42,11 @@ def run(self): class Socks4Code(enum.IntEnum): - success = 0x5a - fail = 0x5b + success = 0x5A + fail = 0x5B class BaseSocks4(AbstractSocks): - @property def version(self): return 4 @@ -61,7 +58,6 @@ def version(self): class Socks4Server(BaseSocks4): - def write_response(self, code): yield from self.io.write_struct(self.fmt, 0, code, 0, self.this_network.packed) @@ -87,7 +83,6 @@ def run(self): class Socks4Client(BaseSocks4): - def resolve_host(self, host): with contextlib.suppress(ValueError): return IPv4Address(host) @@ -109,7 +104,7 @@ class Socks5AuthMethod(enum.IntEnum): no_auth = 0x00 gssapi = 0x01 username_password = 0x02 - no_acceptable = 0xff + no_acceptable = 0xFF class Socks5AddressType(enum.IntEnum): @@ -131,7 +126,6 @@ class Socks5Code(enum.IntEnum): class BaseSocks5(AbstractSocks): - @property def version(self): return 5 @@ -154,7 +148,7 @@ def read_command(self): octets = yield from self.io.read_struct("16s") host = IPv6Address(octets).compressed elif address_type == Socks5AddressType.domain: - host = (yield from self.io.read_pascal_string()) + host = yield from self.io.read_pascal_string() else: raise SocksException(f"Unknown address type {_hex(address_type)}") port = yield from self.io.read_struct("H") @@ -173,7 +167,6 @@ def write_command(self, command, host="0.0.0.0", port=0): class Socks5Server(BaseSocks5): - def auth(self, username, password): auth_methods_count = yield from self.io.read_struct("B") auth_methods = yield from self.io.read_exactly(auth_methods_count) @@ -218,7 +211,6 @@ def run(self, username=None, password=None): class Socks5Client(BaseSocks5): - def auth(self, username, password): auth_required = username is not None if auth_required: @@ -250,13 +242,16 @@ def run(self, host, port, username=None, password=None): yield from self.io.passthrough() -def SocksServer(*, allowed_versions={4, 5}, username=None, password=None, - strict_security_policy=True, encoding=DEFAULT_ENCODING): +def SocksServer( + *, allowed_versions={4, 5}, username=None, password=None, strict_security_policy=True, encoding=DEFAULT_ENCODING +): auth_required = username is not None if 4 in allowed_versions and auth_required and strict_security_policy: - raise SocksException("Socks4 do not provide auth methods, " - "but socks4 allowed and auth provided and " - "strict security policy enabled") + raise SocksException( + "Socks4 do not provide auth methods, " + "but socks4 allowed and auth provided and " + "strict security policy enabled", + ) io = SansIORW(encoding) version = yield from io.read_struct("B", put_back=True) if version not in allowed_versions: @@ -269,8 +264,9 @@ def SocksServer(*, allowed_versions={4, 5}, username=None, password=None, raise SocksException(f"Version {version} is not supported") -def SocksClient(host, port, version, *, username=None, password=None, encoding=DEFAULT_ENCODING, - socks4_extras={}, socks5_extras={}): +def SocksClient( + host, port, version, *, username=None, password=None, encoding=DEFAULT_ENCODING, socks4_extras={}, socks5_extras={} +): auth_required = username is not None if version == 4 and auth_required: raise SocksException("Socks4 do not provide auth methods, but auth provided") diff --git a/siosocks/sansio.py b/src/siosocks/sansio.py similarity index 99% rename from siosocks/sansio.py rename to src/siosocks/sansio.py index 99a5dd1..c768799 100644 --- a/siosocks/sansio.py +++ b/src/siosocks/sansio.py @@ -2,12 +2,10 @@ from .exceptions import SocksException - -MAX_STRING_SIZE = 2 ** 10 +MAX_STRING_SIZE = 2**10 class SansIORW: - def __init__(self, encoding): self.buffer = b"" self.encoding = encoding diff --git a/tests/test_asyncio.py b/tests/test_asyncio.py index 20558bb..c38af9b 100644 --- a/tests/test_asyncio.py +++ b/tests/test_asyncio.py @@ -3,9 +3,8 @@ import pytest import pytest_asyncio -from siosocks.io.asyncio import socks_server_handler, open_connection from siosocks.exceptions import SocksException - +from siosocks.io.asyncio import open_connection, socks_server_handler HOST = "127.0.0.1" MESSAGE = b"socks work!" @@ -46,8 +45,13 @@ async def test_connection_direct_success(endpoint_port): @pytest.mark.asyncio async def test_connection_socks_success(endpoint_port, socks_server_port): - r, w = await open_connection(HOST, endpoint_port, - socks_host=HOST, socks_port=socks_server_port, socks_version=4) + r, w = await open_connection( + HOST, + endpoint_port, + socks_host=HOST, + socks_port=socks_server_port, + socks_version=4, + ) w.write(MESSAGE) m = await r.read(8192) assert m == MESSAGE @@ -56,12 +60,21 @@ async def test_connection_socks_success(endpoint_port, socks_server_port): @pytest.mark.asyncio async def test_connection_socks_failed(socks_server_port, unused_tcp_port): with pytest.raises(SocksException): - await open_connection(HOST, unused_tcp_port, - socks_host=HOST, socks_port=socks_server_port, socks_version=4) + await open_connection( + HOST, + unused_tcp_port, + socks_host=HOST, + socks_port=socks_server_port, + socks_version=4, + ) @pytest.mark.asyncio async def test_connection_partly_passed_error(endpoint_port, socks_server_port): with pytest.raises(SocksException): - await open_connection(HOST, endpoint_port, - socks_host=HOST, socks_port=socks_server_port) + await open_connection( + HOST, + endpoint_port, + socks_host=HOST, + socks_port=socks_server_port, + ) diff --git a/tests/test_protocol.py b/tests/test_protocol.py index 63540d7..d75079f 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -13,7 +13,6 @@ class ConnectionFailed(Exception): class Node: - def __init__(self, generator): self.generator = generator self.receive = [None] @@ -58,7 +57,6 @@ def rotor(client, server, *, fail_connection=False): def test_client_bad_socks_version(): - def server(): io = SansIORW(encoding="utf-8") yield from io.read_exactly(1) @@ -68,7 +66,6 @@ def server(): def test_client_socks4_and_auth(): - def server(): io = SansIORW(encoding="utf-8") yield from io.read_exactly(1) @@ -78,7 +75,6 @@ def server(): def test_client_socks4_connection_failed(): - def server(): io = SansIORW(encoding="utf-8") version, command, port, ipv4 = yield from io.read_struct("BBH4s") @@ -89,7 +85,7 @@ def server(): user_id = yield from io.read_c_string() assert user_id == "yoba" yield from io.connect(ipv4, port) - yield from io.write_struct("BBH4s", 0, 0x5b, 0, b"\x00" * 4) + yield from io.write_struct("BBH4s", 0, 0x5B, 0, b"\x00" * 4) yield from io.passthrough() with pytest.raises(SocksException): @@ -98,7 +94,6 @@ def server(): @pytest.mark.skip("this check removed") def test_client_socks4_redirect_not_supported_by_port(): - def server(): io = SansIORW(encoding="utf-8") version, command, port, ipv4 = yield from io.read_struct("BBH4s") @@ -109,7 +104,7 @@ def server(): user_id = yield from io.read_c_string() assert user_id == "yoba" yield from io.connect(ipv4, port) - yield from io.write_struct("BBH4s", 0, 0x5a, 666, b"\x00" * 4) + yield from io.write_struct("BBH4s", 0, 0x5A, 666, b"\x00" * 4) yield from io.passthrough() with pytest.raises(SocksException): @@ -118,7 +113,6 @@ def server(): @pytest.mark.skip("this check removed") def test_client_socks4_redirect_not_supported_by_host(): - def server(): io = SansIORW(encoding="utf-8") version, command, port, ipv4 = yield from io.read_struct("BBH4s") @@ -129,7 +123,7 @@ def server(): user_id = yield from io.read_c_string() assert user_id == "yoba" yield from io.connect(ipv4, port) - yield from io.write_struct("BBH4s", 0, 0x5a, 0, b"\x7f\x00\x00\x01") + yield from io.write_struct("BBH4s", 0, 0x5A, 0, b"\x7f\x00\x00\x01") yield from io.passthrough() with pytest.raises(SocksException): @@ -137,7 +131,6 @@ def server(): def test_client_socks4_success_by_ipv4(): - def server(): io = SansIORW(encoding="utf-8") version, command, port, ipv4 = yield from io.read_struct("BBH4s") @@ -148,14 +141,13 @@ def server(): user_id = yield from io.read_c_string() assert user_id == "" yield from io.connect(ipv4, port) - yield from io.write_struct("BBH4s", 0, 0x5a, 0, b"\x00" * 4) + yield from io.write_struct("BBH4s", 0, 0x5A, 0, b"\x00" * 4) yield from io.passthrough() rotor(SocksClient("127.0.0.1", 123, 4), server()) def test_client_socks4_success_by_host(): - def server(): io = SansIORW(encoding="utf-8") version, command, port, ipv4 = yield from io.read_struct("BBH4s") @@ -168,14 +160,13 @@ def server(): host = yield from io.read_c_string() assert host == "python.org" yield from io.connect(host, port) - yield from io.write_struct("BBH4s", 0, 0x5a, 0, b"\x00" * 4) + yield from io.write_struct("BBH4s", 0, 0x5A, 0, b"\x00" * 4) yield from io.passthrough() rotor(SocksClient("python.org", 123, 4), server()) def test_server_socks4_auth_required(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write(b"\x04") @@ -186,7 +177,6 @@ def client(): def test_server_socks_bad_socks_version(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write(b"\x06") @@ -197,7 +187,6 @@ def client(): def test_server_socks_bad_socks_version_but_allowed(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write(b"\x06") @@ -208,7 +197,6 @@ def client(): def test_server_socks4_unsupported_command(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("BBH4s", 4, 2, 123, b"\x7f\x00\x00\x01") @@ -220,14 +208,13 @@ def client(): def test_server_socks4_connect_failed(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("BBH4s", 4, 1, 123, b"\x7f\x00\x00\x01") yield from io.write_c_string("yoba") prefix, code, port, ipv4 = yield from io.read_struct("BBH4s") assert prefix == 0 - assert code == 0x5b + assert code == 0x5B assert port == 0 assert ipv4 == b"\x00" * 4 raise RuntimeError("connection failed") @@ -237,14 +224,13 @@ def client(): def test_server_socks4_success_by_ipv4(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("BBH4s", 4, 1, 123, b"\x7f\x00\x00\x01") yield from io.write_c_string("yoba") prefix, code, port, ipv4 = yield from io.read_struct("BBH4s") assert prefix == 0 - assert code == 0x5a + assert code == 0x5A assert port == 0 assert ipv4 == b"\x00" * 4 yield from io.passthrough() @@ -253,7 +239,6 @@ def client(): def test_server_socks4_success_by_host(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("BBH4s", 4, 1, 123, b"\x00\x00\x00\x01") @@ -261,7 +246,7 @@ def client(): yield from io.write_c_string("python.org") prefix, code, port, ipv4 = yield from io.read_struct("BBH4s") assert prefix == 0 - assert code == 0x5a + assert code == 0x5A assert port == 0 assert ipv4 == b"\x00" * 4 yield from io.passthrough() @@ -270,7 +255,6 @@ def client(): def test_client_socks5_request_auth_bad_version(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -285,7 +269,6 @@ def server(): def test_client_socks5_request_auth_not_accepted(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -300,7 +283,6 @@ def server(): def test_client_socks5_request_auth_username_bad_auth_version(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -322,7 +304,6 @@ def server(): def test_client_socks5_request_auth_username_failed(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -344,7 +325,6 @@ def server(): def test_client_socks5_command_bad_version(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -362,7 +342,6 @@ def server(): def test_client_socks5_command_request_not_granted(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -383,7 +362,6 @@ def server(): @pytest.mark.skip("this check removed") def test_client_socks5_command_redirect_is_not_allowed(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -403,7 +381,6 @@ def server(): def test_client_socks5_success_ipv4(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -422,7 +399,6 @@ def server(): def test_client_socks5_success_ipv6(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -441,7 +417,6 @@ def server(): def test_client_socks5_success_domain(): - def server(): io = SansIORW(encoding="utf-8") version, one, auth_method = yield from io.read_struct("BBB") @@ -461,7 +436,6 @@ def server(): def test_server_socks5_no_auth_methods(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -473,7 +447,6 @@ def client(): def test_server_socks5_bad_username_auth_version(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -488,7 +461,6 @@ def client(): def test_server_socks5_bad_username(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -507,7 +479,6 @@ def client(): def test_server_socks5_bad_password(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -526,7 +497,6 @@ def client(): def test_server_socks5_command_not_supported(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -546,7 +516,6 @@ def client(): def test_server_socks5_address_type_not_supported(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -561,7 +530,6 @@ def client(): def test_server_socks5_connection_failed(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -579,7 +547,6 @@ def client(): def test_server_socks5_connection_ipv4_success(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -598,7 +565,6 @@ def client(): def test_server_socks5_connection_ipv6_success(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) @@ -617,7 +583,6 @@ def client(): def test_server_socks5_connection_domain_success(): - def client(): io = SansIORW(encoding="utf-8") yield from io.write_struct("B", 5) diff --git a/tests/test_sansio.py b/tests/test_sansio.py index 77ec360..6a4829a 100644 --- a/tests/test_sansio.py +++ b/tests/test_sansio.py @@ -1,14 +1,12 @@ import pytest -from siosocks.sansio import SansIORW from siosocks.exceptions import SocksException +from siosocks.sansio import SansIORW @pytest.fixture def io(): - class IO: - def __init__(self): self.buffer = [] self.io = SansIORW(encoding="utf-8") @@ -26,7 +24,6 @@ def write(self, data): self.buffer.append(data) def __getattr__(self, name): - def wrapper(*args, **kwargs): gen = getattr(self.io, name)(*args, **kwargs) method, data = gen.send, None diff --git a/tests/test_socketserver.py b/tests/test_socketserver.py index 67c7035..3f70b46 100644 --- a/tests/test_socketserver.py +++ b/tests/test_socketserver.py @@ -6,10 +6,9 @@ import pytest import pytest_asyncio -from siosocks.io.socket import socks_server_handler -from siosocks.io.asyncio import open_connection from siosocks.exceptions import SocksException - +from siosocks.io.asyncio import open_connection +from siosocks.io.socket import socks_server_handler HOST = "127.0.0.1" MESSAGE = b"socks work!" @@ -54,8 +53,13 @@ async def test_connection_direct_success(endpoint_port): @pytest.mark.asyncio async def test_connection_socks_success(endpoint_port, socks_server_port): - r, w = await open_connection(HOST, endpoint_port, - socks_host=HOST, socks_port=socks_server_port, socks_version=4) + r, w = await open_connection( + HOST, + endpoint_port, + socks_host=HOST, + socks_port=socks_server_port, + socks_version=4, + ) w.write(MESSAGE) m = await r.read(8192) assert m == MESSAGE @@ -64,5 +68,9 @@ async def test_connection_socks_success(endpoint_port, socks_server_port): @pytest.mark.asyncio async def test_connection_partly_passed_error(endpoint_port, socks_server_port): with pytest.raises(SocksException): - await open_connection(HOST, endpoint_port, - socks_host=HOST, socks_port=socks_server_port) + await open_connection( + HOST, + endpoint_port, + socks_host=HOST, + socks_port=socks_server_port, + ) diff --git a/tests/test_trio.py b/tests/test_trio.py index d6fa680..d96d159 100644 --- a/tests/test_trio.py +++ b/tests/test_trio.py @@ -1,11 +1,10 @@ from functools import partial -import trio import pytest +import trio from siosocks.exceptions import SocksException -from siosocks.io.trio import socks_server_handler, open_tcp_stream - +from siosocks.io.trio import open_tcp_stream, socks_server_handler # TODO: Use fixtures after https://github.com/pytest-dev/pytest-asyncio/issues/124 resolved @@ -14,7 +13,6 @@ async def endpoint(nursery): - async def handler(stream): async with stream: while True: @@ -48,8 +46,13 @@ async def test_connection_direct_success(nursery): async def test_connection_socks_success(nursery): endpoint_port = await endpoint(nursery) socks_server_port = await socks(nursery) - stream = await open_tcp_stream(HOST, endpoint_port, - socks_host=HOST, socks_port=socks_server_port, socks_version=4) + stream = await open_tcp_stream( + HOST, + endpoint_port, + socks_host=HOST, + socks_port=socks_server_port, + socks_version=4, + ) async with stream: await stream.send_all(MESSAGE) m = await stream.receive_some(8192) @@ -61,5 +64,9 @@ async def test_connection_partly_passed_error(nursery): endpoint_port = await endpoint(nursery) socks_server_port = await socks(nursery) with pytest.raises(SocksException): - await open_tcp_stream(HOST, endpoint_port, - socks_host=HOST, socks_port=socks_server_port) + await open_tcp_stream( + HOST, + endpoint_port, + socks_host=HOST, + socks_port=socks_server_port, + )