├── .coveragerc ├── .github └── workflows │ └── pythonpackage.yml ├── .gitignore ├── .readthedocs.yml ├── .travis.yml ├── LICENSE ├── README.md ├── _config.yml ├── docs ├── Makefile ├── conf.py ├── index.rst ├── iofree.exceptions.rst ├── iofree.rst ├── iofree.schema.rst ├── make.bat ├── modules.rst └── requirements.txt ├── iofree ├── __init__.py ├── contrib │ ├── common.py │ └── socks5.py ├── exceptions.py └── schema.py ├── pyproject.toml ├── setup.cfg ├── setup.py └── tests ├── test_iofree.py ├── test_parser.py ├── test_schema.py └── test_socks5.py /.coveragerc: -------------------------------------------------------------------------------- 1 | [run] 2 | source = iofree 3 | 4 | [report] 5 | include = 6 | iofree/*.py 7 | tests/*.py 8 | -------------------------------------------------------------------------------- /.github/workflows/pythonpackage.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests and lint with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: Python package 5 | 6 | on: 7 | push: 8 | branches: [ master ] 9 | pull_request: 10 | branches: [ master ] 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | strategy: 17 | matrix: 18 | python-version: [3.6, 3.7, 3.8] 19 | 20 | steps: 21 | - uses: actions/checkout@v2 22 | - name: Set up Python ${{ matrix.python-version }} 23 | uses: actions/setup-python@v1 24 | with: 25 | python-version: ${{ matrix.python-version }} 26 | - name: Install dependencies 27 | run: | 28 | python -m pip install --upgrade pip 29 | - name: Lint with flake8 30 | run: | 31 | pip install flake8 32 | # stop the build if there are Python syntax errors or undefined names 33 | flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics 34 | # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide 35 | flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics 36 | - name: Black Code Formatter 37 | uses: lgeiger/black-action@v1.0.1 38 | - name: Test with pytest 39 | run: | 40 | pip install pytest 41 | python setup.py test 42 | - name: Upload coverage to Codecov 43 | uses: codecov/codecov-action@v1 44 | with: 45 | file: ./coverage.xml 46 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | -------------------------------------------------------------------------------- /.readthedocs.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | sphinx: 3 | configuration: docs/conf.py 4 | formats: all 5 | python: 6 | version: 3.7 7 | install: 8 | - requirements: docs/requirements.txt 9 | - method: setuptools 10 | path: . 11 | # system_packages: true 12 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | dist: xenial 3 | sudo: required 4 | python: 5 | - 3.6 6 | - 3.7 7 | install: 8 | - python setup.py install 9 | - pip install codecov 10 | script: 11 | - python setup.py test 12 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Yingbo Gu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # iofree 2 | 3 | [![Build Status](https://travis-ci.org/guyingbo/iofree.svg?branch=master)](https://travis-ci.org/guyingbo/iofree) 4 | [![Documentation Status](https://readthedocs.org/projects/iofree/badge/?version=latest)](https://iofree.readthedocs.io/en/latest/?badge=latest) 5 | [![Python Version](https://img.shields.io/pypi/pyversions/iofree.svg)](https://pypi.python.org/pypi/iofree) 6 | [![Version](https://img.shields.io/pypi/v/iofree.svg)](https://pypi.python.org/pypi/iofree) 7 | [![Format](https://img.shields.io/pypi/format/iofree.svg)](https://pypi.python.org/pypi/iofree) 8 | [![License](https://img.shields.io/pypi/l/iofree.svg)](https://pypi.python.org/pypi/iofree) 9 | [![codecov](https://codecov.io/gh/guyingbo/iofree/branch/master/graph/badge.svg)](https://codecov.io/gh/guyingbo/iofree) 10 | 11 | `iofree` is an easy-to-use and powerful library to help you implement network protocols and binary parsers. 12 | 13 | ## Installation 14 | 15 | ~~~ 16 | pip install iofree 17 | ~~~ 18 | 19 | ## Advantages 20 | 21 | Using iofree, you can: 22 | 23 | * define network protocols and file format in a clear and precise manner 24 | * parse both binary streams and files 25 | 26 | ## Documentation 27 | 28 | ### Basic Usage 29 | 30 | ```python 31 | >>> from iofree import schema 32 | >>> schema.uint8(1) 33 | b'\x01' 34 | >>> schema.uint32be(3) 35 | b'\x00\x00\x00\x03' 36 | >>> schema.uint32be.parse(b'\x00\x00\x00\x03') 37 | 3 38 | ``` 39 | 40 | ### Tutorial 1: a simple parser 41 | 42 | ```python 43 | >>> class Simple(schema.BinarySchema): 44 | ... a = schema.uint8 45 | ... b = schema.uint32be # "be" for big-endian 46 | ... 47 | >>> Simple(1, 3).binary 48 | b'\x01\x00\x00\x00\x03' 49 | >>> binary = _ 50 | >>> Simple.parse(binary) 51 | 52 | ``` 53 | 54 | ### Built-in units: 55 | 56 | commonly used number units: 57 | 58 | * int8 uint8 59 | * int16 int16be uint16 uint16be 60 | * int24 int24be uint24 uint24be 61 | * int32 int32be uint32 uint32be 62 | * int64 int64be uint64 uint64be 63 | * float16 float16be 64 | * float32 float32be 65 | * float64 float64be 66 | 67 | simple units: 68 | 69 | * Bytes 70 | * String 71 | * EndWith 72 | 73 | composite units: 74 | 75 | * LengthPrefixedBytes 76 | * LengthPrefixedString 77 | * LengthPrefixedObjectList 78 | * LengthPrefixedObject 79 | * MustEqual 80 | * Switch 81 | * SizedIntEnum 82 | * Convert 83 | * Group 84 | 85 | [API docs](https://iofree.readthedocs.io/en/latest/index.html) 86 | 87 | Here is a real life example [definition](https://github.com/guyingbo/iofree/blob/master/iofree/contrib/socks5.py) of socks5 client request, you can see the following code snippet: 88 | 89 | ```python 90 | class Socks5ClientRequest(schema.BinarySchema): 91 | ver = schema.MustEqual(schema.uint8, 5) 92 | cmd = schema.SizedIntEnum(schema.uint8, Cmd) 93 | rsv = schema.MustEqual(schema.uint8, 0) 94 | addr = Addr 95 | ``` 96 | 97 | ### Tutorial 2: define socks5 address format 98 | 99 | ```python 100 | In [1]: import socket 101 | ...: from iofree import schema 102 | ...: 103 | ...: 104 | ...: class Addr(schema.BinarySchema): 105 | ...: atyp: int = schema.uint8 106 | ...: host: str = schema.Switch( 107 | ...: "atyp", 108 | ...: { 109 | ...: 1: schema.Convert( 110 | ...: schema.Bytes(4), encode=socket.inet_aton, decode=socket.inet_ntoa 111 | ...: 112 | ...: ), 113 | ...: 4: schema.Convert( 114 | ...: schema.Bytes(16), 115 | ...: encode=lambda x: socket.inet_pton(socket.AF_INET6, x), 116 | ...: decode=lambda x: socket.inet_ntop(socket.AF_INET6, x), 117 | ...: ), 118 | ...: 3: schema.LengthPrefixedString(schema.uint8), 119 | ...: }, 120 | ...: ) 121 | ...: port: int = schema.uint16be 122 | ...: 123 | 124 | In [2]: addr = Addr(1, '172.16.1.20', 80) 125 | 126 | In [3]: addr 127 | Out[3]: 128 | 129 | In [4]: addr.binary 130 | Out[4]: b'\x01\xac\x10\x01\x14\x00P' 131 | 132 | In [5]: Addr.parse(addr.binary) 133 | Out[5]: 134 | ``` 135 | 136 | A complete socks5 Addr [definition](https://github.com/guyingbo/iofree/blob/master/iofree/contrib/common.py) 137 | 138 | ## Projects using iofree 139 | 140 | * [Shadowproxy](https://github.com/guyingbo/shadowproxy) 141 | * [socks5 models](https://github.com/guyingbo/iofree/blob/master/iofree/contrib/socks5.py) and [socks5 protocol](https://github.com/guyingbo/shadowproxy/blob/master/shadowproxy/protocols/socks5.py) 142 | * [shadowsocks parser](https://github.com/guyingbo/shadowproxy/blob/master/shadowproxy/proxies/shadowsocks/parser.py) 143 | * [shadowsocks aead parser](https://github.com/guyingbo/shadowproxy/blob/master/shadowproxy/proxies/aead/parser.py) 144 | * [python tls1.3](https://github.com/guyingbo/tls1.3) 145 | * [TLS1.3 models](https://github.com/guyingbo/tls1.3/blob/master/tls/models.py) and [protocol](https://github.com/guyingbo/tls1.3/blob/master/tls/session.py) 146 | 147 | ## References 148 | 149 | `iofree` parser is inspired by project [ohneio](https://github.com/acatton/ohneio). 150 | -------------------------------------------------------------------------------- /_config.yml: -------------------------------------------------------------------------------- 1 | theme: jekyll-theme-cayman -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = . 9 | BUILDDIR = _build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | # import os 14 | # import sys 15 | # sys.path.insert(0, '/Users/mac/Projects/iofree/iofree') 16 | 17 | 18 | # -- Project information ----------------------------------------------------- 19 | 20 | project = 'iofree' 21 | copyright = '2020, GuYingbo' 22 | author = 'GuYingbo' 23 | 24 | 25 | # -- General configuration --------------------------------------------------- 26 | 27 | # Add any Sphinx extension module names here, as strings. They can be 28 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 29 | # ones. 30 | extensions = [ 31 | 'sphinx.ext.autodoc', 32 | 'sphinx.ext.viewcode', 33 | 'sphinx.ext.todo', 34 | 'sphinx.ext.autodoc', 35 | ] 36 | 37 | # Add any paths that contain templates here, relative to this directory. 38 | templates_path = ['_templates'] 39 | 40 | # The language for content autogenerated by Sphinx. Refer to documentation 41 | # for a list of supported languages. 42 | # 43 | # This is also used if you do content translation via gettext catalogs. 44 | # Usually you set "language" from the command line for these cases. 45 | language = 'en' 46 | 47 | # List of patterns, relative to source directory, that match files and 48 | # directories to ignore when looking for source files. 49 | # This pattern also affects html_static_path and html_extra_path. 50 | exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] 51 | 52 | 53 | # -- Options for HTML output ------------------------------------------------- 54 | 55 | # The theme to use for HTML and HTML Help pages. See the documentation for 56 | # a list of builtin themes. 57 | # 58 | html_theme = 'alabaster' 59 | 60 | # Add any paths that contain custom static files (such as style sheets) here, 61 | # relative to this directory. They are copied after the builtin static files, 62 | # so a file named "default.css" will overwrite the builtin "default.css". 63 | html_static_path = ['_static'] 64 | 65 | 66 | # -- Extension configuration ------------------------------------------------- 67 | 68 | # -- Options for todo extension ---------------------------------------------- 69 | 70 | # If true, `todo` and `todoList` produce output, else they produce nothing. 71 | todo_include_todos = True -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- 1 | .. iofree documentation master file, created by 2 | sphinx-quickstart on Tue Apr 7 01:16:51 2020. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | Welcome to iofree's documentation! 7 | ================================== 8 | 9 | .. toctree:: 10 | :maxdepth: 4 11 | :caption: Contents: 12 | 13 | iofree 14 | 15 | 16 | Indices and tables 17 | ================== 18 | 19 | * :ref:`genindex` 20 | * :ref:`modindex` 21 | * :ref:`search` 22 | -------------------------------------------------------------------------------- /docs/iofree.exceptions.rst: -------------------------------------------------------------------------------- 1 | iofree.exceptions module 2 | ======================== 3 | 4 | .. automodule:: iofree.exceptions 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/iofree.rst: -------------------------------------------------------------------------------- 1 | iofree package 2 | ============== 3 | 4 | .. automodule:: iofree 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | 9 | Submodules 10 | ---------- 11 | 12 | .. toctree:: 13 | :maxdepth: 4 14 | 15 | iofree.exceptions 16 | iofree.schema 17 | -------------------------------------------------------------------------------- /docs/iofree.schema.rst: -------------------------------------------------------------------------------- 1 | iofree.schema module 2 | ==================== 3 | 4 | .. automodule:: iofree.schema 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=. 11 | set BUILDDIR=_build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/modules.rst: -------------------------------------------------------------------------------- 1 | iofree 2 | ====== 3 | 4 | .. toctree:: 5 | :maxdepth: 4 6 | 7 | iofree 8 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | Sphinx==3.0.0 2 | Pygments==2.7.4 3 | setuptools==70.0.0 4 | docutils==0.14 5 | mock==1.0.1 6 | pillow>=6.2.2 7 | alabaster>=0.7,<0.8,!=0.7.5 8 | commonmark==0.8.1 9 | recommonmark==0.5.0 10 | sphinx-rtd-theme<0.5 11 | readthedocs-sphinx-ext<1.1 12 | -------------------------------------------------------------------------------- /iofree/__init__.py: -------------------------------------------------------------------------------- 1 | """`iofree` is an easy-to-use and powerful library \ 2 | to help you implement network protocols and binary parsers.""" 3 | import sys 4 | import typing 5 | from collections import deque 6 | from enum import IntEnum, auto 7 | from socket import SocketType 8 | from struct import Struct 9 | 10 | from .exceptions import NoResult, ParseError 11 | 12 | __version__ = "0.2.4" 13 | _wait = object() 14 | _no_result = object() 15 | 16 | 17 | class Traps(IntEnum): 18 | _read = auto() 19 | _read_more = auto() 20 | _read_until = auto() 21 | _read_struct = auto() 22 | _read_int = auto() 23 | _wait = auto() 24 | _peek = auto() 25 | _wait_event = auto() 26 | _get_parser = auto() 27 | 28 | 29 | class State(IntEnum): 30 | _state_wait = auto() 31 | _state_next = auto() 32 | _state_end = auto() 33 | 34 | 35 | class Parser: 36 | def __init__(self, gen: typing.Generator): 37 | self.gen = gen 38 | self._input = bytearray() 39 | self._input_events: typing.Deque = deque() 40 | self._output_events: typing.Deque = deque() 41 | self._res = _no_result 42 | self._mapping_stack: typing.Deque = deque() 43 | self._next_value = None 44 | self._last_trap: typing.Optional[tuple] = None 45 | self._pos = 0 46 | self._state: State = State._state_wait 47 | self._process() 48 | 49 | def __repr__(self): 50 | return f"<{self.__class__.__qualname__}({self.gen})>" 51 | 52 | def __iter__(self): 53 | return self 54 | 55 | def __next__(self) -> typing.Any: 56 | if self._output_events: 57 | return self._output_events.popleft() 58 | raise StopIteration 59 | 60 | def parse(self, data: bytes, *, strict: bool = True) -> typing.Any: 61 | """ 62 | parse bytes 63 | """ 64 | self.send(data) 65 | if strict and self.has_more_data(): 66 | raise ParseError("redundant data left") 67 | return self.get_result() 68 | 69 | def send(self, data: bytes = b"") -> None: 70 | """ 71 | send data for parsing 72 | """ 73 | self._input.extend(data) 74 | self._process() 75 | 76 | def read_output_bytes(self) -> bytes: 77 | buf = [] 78 | for to_send, close, exc, result in self: 79 | buf.append(result) 80 | return b"".join(buf) 81 | 82 | def respond( 83 | self, 84 | *, 85 | data: bytes = b"", 86 | close: bool = False, 87 | exc: typing.Optional[Exception] = None, 88 | result: typing.Any = _no_result, 89 | ) -> None: 90 | """produce some event data to interact with a stream: 91 | data: bytes to send to the peer 92 | close: whether the socket should be closed 93 | exc: raise an exception to break the loop 94 | result: result to return 95 | """ 96 | self._output_events.append((data, close, exc, result)) 97 | 98 | def run(self, sock: SocketType) -> typing.Any: 99 | "reference implementation of how to deal with socket" 100 | self.send(b"") 101 | while True: 102 | for to_send, close, exc, result in self: 103 | if to_send: 104 | sock.sendall(to_send) 105 | if close: 106 | sock.close() 107 | if exc: 108 | raise exc 109 | if result is not _no_result: 110 | return result 111 | data = sock.recv(1024) 112 | if not data: 113 | raise ParseError("need data") 114 | self.send(data) 115 | 116 | @property 117 | def has_result(self) -> bool: 118 | return self._res is not _no_result 119 | 120 | def get_result(self) -> typing.Any: 121 | """ 122 | raises *NoResult* exception if no result has been set 123 | """ 124 | self._process() 125 | if not self.has_result: 126 | raise NoResult("no result") 127 | return self._res 128 | 129 | def set_result(self, result) -> None: 130 | self._res = result 131 | self.respond(result=result) 132 | 133 | def finished(self) -> bool: 134 | return self._state is State._state_end 135 | 136 | def _process(self) -> None: 137 | if self._state is State._state_end: 138 | return 139 | self._state = State._state_next 140 | while self._state is State._state_next: 141 | self._next_state() 142 | 143 | def _next_state(self) -> None: 144 | if self._last_trap is None: 145 | try: 146 | trap, *args = self.gen.send(self._next_value) 147 | except StopIteration as e: 148 | self._state = State._state_end 149 | self.set_result(e.value) 150 | return 151 | except Exception: 152 | self._state = State._state_end 153 | tb = sys.exc_info()[2] 154 | raise ParseError(f"{self._next_value!r}").with_traceback(tb) 155 | else: 156 | if not isinstance(trap, Traps): 157 | self._state = State._state_end 158 | raise RuntimeError(f"Expect Traps object, but got: {trap}") 159 | else: 160 | trap, *args = self._last_trap 161 | result = getattr(self, trap.name)(*args) 162 | if result is _wait: 163 | self._state = State._state_wait 164 | self._last_trap = (trap, *args) 165 | else: 166 | self._state = State._state_next 167 | self._next_value = result 168 | self._last_trap = None 169 | 170 | def readall(self) -> bytes: 171 | """ 172 | retrieve data from input back 173 | """ 174 | return self._read(0) 175 | 176 | def has_more_data(self) -> bool: 177 | "indicate whether input has some bytes left" 178 | return len(self._input) > 0 179 | 180 | def send_event(self, event: typing.Any) -> None: 181 | self._input_events.append(event) 182 | self._process() 183 | 184 | def _wait_event(self): 185 | if self._input_events: 186 | return self._input_events.popleft() 187 | return _wait 188 | 189 | def _wait(self) -> typing.Optional[object]: 190 | if not getattr(self, "_waiting", False): 191 | self._waiting = True 192 | return _wait 193 | self._waiting = False 194 | return None 195 | 196 | def _read(self, nbytes: int = 0, from_=None) -> bytes: 197 | buf = self._input if from_ is None else from_ 198 | if nbytes == 0: 199 | data = bytes(buf) 200 | del buf[:] 201 | return data 202 | if len(buf) < nbytes: 203 | return _wait 204 | data = bytes(buf[:nbytes]) 205 | del buf[:nbytes] 206 | return data 207 | 208 | def _read_more(self, nbytes: int = 1, from_=None) -> typing.Union[object, bytes]: 209 | buf = self._input if from_ is None else from_ 210 | if len(buf) < nbytes: 211 | return _wait 212 | data = bytes(buf) 213 | del buf[:] 214 | return data 215 | 216 | def _read_until( 217 | self, data: bytes, return_tail: bool = True, from_=None 218 | ) -> typing.Union[object, bytes]: 219 | buf = self._input if from_ is None else from_ 220 | index = buf.find(data, self._pos) 221 | if index == -1: 222 | self._pos = len(buf) - len(data) + 1 223 | self._pos = self._pos if self._pos > 0 else 0 224 | return _wait 225 | size = index + len(data) 226 | if return_tail: 227 | data = bytes(buf[:size]) 228 | else: 229 | data = bytes(buf[:index]) 230 | del buf[:size] 231 | self._pos = 0 232 | return data 233 | 234 | def _read_struct( 235 | self, struct_obj: Struct, from_=None 236 | ) -> typing.Union[object, tuple]: 237 | buf = self._input if from_ is None else from_ 238 | size = struct_obj.size 239 | if len(buf) < size: 240 | return _wait 241 | result = struct_obj.unpack_from(buf) 242 | del buf[:size] 243 | return result 244 | 245 | def _read_int( 246 | self, nbytes: int, byteorder: str = "big", signed: bool = False, from_=None 247 | ) -> typing.Union[object, int]: 248 | buf = self._input if from_ is None else from_ 249 | if len(buf) < nbytes: 250 | return _wait 251 | data = self._read(nbytes) 252 | return int.from_bytes(data, byteorder, signed=signed) 253 | 254 | def _peek(self, nbytes: int = 1, from_=None) -> typing.Union[object, bytes]: 255 | buf = self._input if from_ is None else from_ 256 | if len(buf) < nbytes: 257 | return _wait 258 | return bytes(buf[:nbytes]) 259 | 260 | def _get_parser(self) -> "Parser": 261 | return self 262 | 263 | 264 | class LinkedNode: 265 | __slots__ = ("parser", "next") 266 | 267 | def __init__(self, parser: Parser, next_: typing.Optional["LinkedNode"]): 268 | self.parser = parser 269 | self.next = next_ 270 | 271 | 272 | class ParserChain: 273 | def __init__(self, *parsers: Parser): 274 | nxt = None 275 | for parser in reversed(parsers): 276 | node = LinkedNode(parser, nxt) 277 | nxt = node 278 | self.first = node 279 | 280 | def send(self, data: bytes) -> None: 281 | self.first.parser.send(data) 282 | 283 | def __iter__(self): 284 | return self._get_events(self.first) 285 | 286 | def _get_events( 287 | self, node: LinkedNode 288 | ) -> typing.Generator[ 289 | typing.Tuple[ 290 | typing.Optional[bytes], 291 | typing.Optional[bool], 292 | typing.Optional[Exception], 293 | typing.Any, 294 | ], 295 | None, 296 | None, 297 | ]: 298 | for data, close, exc, result in node.parser: 299 | if result is not _no_result and node.next: 300 | node.next.parser.send(result) 301 | yield (data, close, exc, _no_result) 302 | else: 303 | yield (data, close, exc, result) 304 | if node.next: 305 | yield from self._get_events(node.next) 306 | 307 | 308 | def read(nbytes: int = 0, *, from_=None) -> typing.Generator[tuple, bytes, bytes]: 309 | """ 310 | if nbytes = 0, read as many as possible, empty bytes is valid; 311 | if nbytes > 0, read *exactly* ``nbytes`` 312 | """ 313 | return (yield (Traps._read, nbytes, from_)) 314 | 315 | 316 | def read_more(nbytes: int = 1, *, from_=None) -> typing.Generator[tuple, bytes, bytes]: 317 | """ 318 | read *at least* ``nbytes`` 319 | """ 320 | if nbytes <= 0: 321 | raise ValueError(f"nbytes must > 0, but got {nbytes}") 322 | return (yield (Traps._read_more, nbytes, from_)) 323 | 324 | 325 | def read_until( 326 | data: bytes, *, return_tail: bool = True, from_=None 327 | ) -> typing.Generator[tuple, bytes, bytes]: 328 | """ 329 | read until some bytes appear 330 | """ 331 | return (yield (Traps._read_until, data, return_tail, from_)) 332 | 333 | 334 | def read_struct(fmt: str, *, from_=None) -> typing.Generator[tuple, tuple, tuple]: 335 | """ 336 | read specific formatted data 337 | """ 338 | return (yield (Traps._read_struct, Struct(fmt), from_)) 339 | 340 | 341 | def read_raw_struct( 342 | struct_obj: Struct, *, from_=None 343 | ) -> typing.Generator[tuple, tuple, tuple]: 344 | """ 345 | read raw struct formatted data 346 | """ 347 | return (yield (Traps._read_struct, struct_obj, from_)) 348 | 349 | 350 | def read_int( 351 | nbytes: int, byteorder: str = "big", *, signed: bool = False, from_=None 352 | ) -> typing.Generator[tuple, int, int]: 353 | """ 354 | read some bytes as integer 355 | """ 356 | if nbytes <= 0: 357 | raise ValueError(f"nbytes must > 0, but got {nbytes}") 358 | return (yield (Traps._read_int, nbytes, byteorder, signed, from_)) 359 | 360 | 361 | def wait() -> typing.Generator[tuple, bytes, typing.Optional[object]]: 362 | """ 363 | wait for next send event 364 | """ 365 | return (yield (Traps._wait,)) 366 | 367 | 368 | def peek(nbytes: int = 1, *, from_=None) -> typing.Generator[tuple, bytes, bytes]: 369 | """ 370 | peek many bytes without taking them away from buffer 371 | """ 372 | if nbytes <= 0: 373 | raise ValueError(f"nbytes must > 0, but got {nbytes}") 374 | return (yield (Traps._peek, nbytes, from_)) 375 | 376 | 377 | def wait_event() -> typing.Generator[tuple, typing.Any, typing.Any]: 378 | """ 379 | wait for an event 380 | """ 381 | return (yield (Traps._wait_event,)) 382 | 383 | 384 | def get_parser() -> typing.Generator[tuple, Parser, Parser]: 385 | "get current parser object" 386 | return (yield (Traps._get_parser,)) 387 | 388 | 389 | def parser(generator_func: typing.Callable) -> typing.Callable: 390 | "decorator function to wrap a generator" 391 | 392 | def create_parser(*args, **kwargs) -> Parser: 393 | return Parser(generator_func(*args, **kwargs)) 394 | 395 | generator_func.parser = create_parser 396 | return generator_func 397 | -------------------------------------------------------------------------------- /iofree/contrib/common.py: -------------------------------------------------------------------------------- 1 | import socket 2 | 3 | from .. import schema 4 | 5 | 6 | class Addr(schema.BinarySchema): 7 | atyp: int = schema.uint8 8 | host: str = schema.Switch( 9 | "atyp", 10 | { 11 | 1: schema.Convert( 12 | schema.Bytes(4), encode=socket.inet_aton, decode=socket.inet_ntoa 13 | ), 14 | 4: schema.Convert( 15 | schema.Bytes(16), 16 | encode=lambda x: socket.inet_pton(socket.AF_INET6, x), 17 | decode=lambda x: socket.inet_ntop(socket.AF_INET6, x), 18 | ), 19 | 3: schema.LengthPrefixedString(schema.uint8), 20 | }, 21 | ) 22 | port: int = schema.uint16be 23 | 24 | @classmethod 25 | def from_tuple(cls, addr): 26 | try: 27 | return cls(1, *addr) 28 | except OSError: 29 | try: 30 | return cls(4, *addr) 31 | except OSError: 32 | return cls(3, *addr) 33 | -------------------------------------------------------------------------------- /iofree/contrib/socks5.py: -------------------------------------------------------------------------------- 1 | # references: 2 | # rfc1928(SOCKS Protocol Version 5): https://www.ietf.org/rfc/rfc1928.txt 3 | # rfc1929(Username/Password Authentication for SOCKS V5): 4 | # https://tools.ietf.org/html/rfc1929 5 | # handshake server selection 6 | # +----+----------+----------+ +----+--------+ 7 | # |VER | NMETHODS | METHODS | |VER | METHOD | 8 | # +----+----------+----------+ +----+--------+ 9 | # | 1 | 1 | 1 to 255 | | 1 | 1 | 10 | # +----+----------+----------+ +----+--------+ 11 | # Username/Password Authentication auth reply 12 | # +----+------+----------+------+----------+ +----+--------+ 13 | # |VER | ULEN | UNAME | PLEN | PASSWD | |VER | STATUS | 14 | # +----+------+----------+------+----------+ +----+--------+ 15 | # | 1 | 1 | 1 to 255 | 1 | 1 to 255 | | 1 | 1 | 16 | # +----+------+----------+------+----------+ +----+--------+ 17 | # request 18 | # +----+-----+-------+------+----------+----------+ 19 | # |VER | CMD | RSV | ATYP | DST.ADDR | DST.PORT | 20 | # +----+-----+-------+------+----------+----------+ 21 | # | 1 | 1 | X'00' | 1 | Variable | 2 | 22 | # +----+-----+-------+------+----------+----------+ 23 | # reply 24 | # +----+-----+-------+------+----------+----------+ 25 | # |VER | REP | RSV | ATYP | BND.ADDR | BND.PORT | 26 | # +----+-----+-------+------+----------+----------+ 27 | # | 1 | 1 | X'00' | 1 | Variable | 2 | 28 | # +----+-----+-------+------+----------+----------+ 29 | # udp relay request and reply 30 | # +----+------+------+----------+----------+----------+ 31 | # |RSV | FRAG | ATYP | DST.ADDR | DST.PORT | DATA | 32 | # +----+------+------+----------+----------+----------+ 33 | # | 2 | 1 | 1 | Variable | 2 | Variable | 34 | # +----+------+------+----------+----------+----------+ 35 | import enum 36 | 37 | from .. import schema 38 | from .common import Addr 39 | 40 | 41 | class AuthMethod(enum.IntEnum): 42 | no_auth = 0 43 | gssapi = 1 44 | user_auth = 2 45 | no_acceptable_method = 255 46 | 47 | 48 | class Cmd(enum.IntEnum): 49 | connect = 1 50 | bind = 2 51 | associate = 3 52 | 53 | 54 | class Rep(enum.IntEnum): 55 | succeeded = 0 56 | general_failure = 1 57 | not_allowed = 2 58 | network_unreachable = 3 59 | host_unreachable = 4 60 | connection_refused = 5 61 | ttl_expired = 6 62 | command_not_supported = 7 63 | address_type_not_supported = 8 64 | 65 | 66 | class Handshake(schema.BinarySchema): 67 | ver = schema.MustEqual(schema.uint8, 5) 68 | methods = schema.LengthPrefixedObjectList( 69 | schema.uint8, schema.SizedIntEnum(schema.uint8, AuthMethod) 70 | ) 71 | 72 | 73 | class ServerSelection(schema.BinarySchema): 74 | ver = schema.MustEqual(schema.uint8, 5) 75 | method = schema.SizedIntEnum(schema.uint8, AuthMethod) 76 | 77 | 78 | class UsernameAuth(schema.BinarySchema): 79 | auth_ver = schema.MustEqual(schema.uint8, 1) 80 | username = schema.LengthPrefixedString(schema.uint8) 81 | password = schema.LengthPrefixedString(schema.uint8) 82 | 83 | 84 | class UsernameAuthReply(schema.BinarySchema): 85 | auth_ver = schema.MustEqual(schema.uint8, 1) 86 | status = schema.MustEqual(schema.uint8, 0) 87 | 88 | 89 | class ClientRequest(schema.BinarySchema): 90 | ver = schema.MustEqual(schema.uint8, 5) 91 | cmd = schema.SizedIntEnum(schema.uint8, Cmd) 92 | rsv = schema.MustEqual(schema.uint8, 0) 93 | addr = Addr 94 | 95 | 96 | class Reply(schema.BinarySchema): 97 | ver = schema.MustEqual(schema.uint8, 5) 98 | rep = schema.SizedIntEnum(schema.uint8, Rep) 99 | rsv = schema.MustEqual(schema.uint8, 0) 100 | bind_addr = Addr 101 | 102 | 103 | class UDPRelay(schema.BinarySchema): 104 | rsv = schema.MustEqual(schema.Bytes(2), b"\x00\x00") 105 | flag = schema.uint8 106 | addr = Addr 107 | data = schema.Bytes(-1) 108 | -------------------------------------------------------------------------------- /iofree/exceptions.py: -------------------------------------------------------------------------------- 1 | class NoResult(Exception): 2 | "" 3 | 4 | 5 | class ParseError(Exception): 6 | "" 7 | -------------------------------------------------------------------------------- /iofree/schema.py: -------------------------------------------------------------------------------- 1 | import abc 2 | import enum 3 | import struct 4 | import typing 5 | from collections import deque 6 | from struct import Struct 7 | 8 | from . import ( 9 | Parser, 10 | get_parser, 11 | read, 12 | read_int, 13 | read_raw_struct, 14 | read_struct, 15 | read_until, 16 | wait, 17 | ) 18 | from .exceptions import ParseError 19 | 20 | _parent_stack: typing.Deque["BinarySchema"] = deque() 21 | 22 | 23 | class Unit(abc.ABC): 24 | """Unit is the base class of all units. \ 25 | If you can build your own unit class, you must inherit from it""" 26 | 27 | def __iter__(self): 28 | return self.get_value() 29 | 30 | @abc.abstractmethod 31 | def get_value(self) -> typing.Generator: 32 | "get object you want from bytes" 33 | 34 | @abc.abstractmethod 35 | def __call__(self, obj: typing.Any) -> bytes: 36 | "convert user-given object to bytes" 37 | 38 | def parse(self, data: bytes, *, strict: bool = True): 39 | "a convenient function to help you parse fixed bytes" 40 | return Parser(self.get_value()).parse(data, strict=strict) 41 | 42 | 43 | class BinarySchemaMetaclass(type): 44 | def __new__(mcls, name, bases, namespace, **kwargs): 45 | fields: typing.Dict[str, FieldType] = {} 46 | for key, member in namespace.items(): 47 | if isinstance(member, (Unit, BinarySchemaMetaclass)): 48 | fields[key] = member 49 | namespace[key] = MemberDescriptor(key, member) 50 | namespace["_fields"] = fields 51 | return super().__new__(mcls, name, bases, namespace) 52 | 53 | # def __init__(cls, name: str, bases: tuple, namespace: dict): 54 | # fields: typing.Dict[str, FieldType] = {} 55 | # for key, member in namespace.items(): 56 | # if isinstance(member, (Unit, BinarySchemaMetaclass)): 57 | # fields[key] = member 58 | # cls._fields = fields 59 | # super().__init__(name, bases, namespace) 60 | 61 | def __str__(cls): 62 | s = ", ".join(f"{name}={field}" for name, field in cls._fields.items()) 63 | return f"{cls.__name__}({s})" 64 | 65 | def __iter__(cls): 66 | return cls.get_value() 67 | 68 | def get_value(cls) -> typing.Generator[tuple, typing.Any, "BinarySchema"]: 69 | "get `BinarySchema` object from bytes" 70 | mapping: typing.Dict[str, typing.Any] = {} 71 | parser = yield from get_parser() 72 | parser._mapping_stack.append(mapping) 73 | try: 74 | for name, field in cls._fields.items(): 75 | mapping[name] = yield from field.get_value() 76 | except Exception: 77 | raise ParseError(mapping) 78 | finally: 79 | parser._mapping_stack.pop() 80 | return cls(*mapping.values()) 81 | 82 | def get_parser(cls) -> Parser: 83 | return Parser(cls.get_value()) 84 | 85 | def parse(cls, data: bytes, *, strict: bool = True) -> "BinarySchema": 86 | return cls.get_parser().parse(data, strict=strict) 87 | 88 | 89 | class BinarySchema(metaclass=BinarySchemaMetaclass): 90 | """The main class for users to define their own binary structures""" 91 | 92 | def __init__(self, *args): 93 | self._modified = True 94 | if len(args) != len(self.__class__._fields): 95 | raise ValueError( 96 | f"need {len(self.__class__._fields)} args, got {len(args)}" 97 | ) 98 | self.values = {} 99 | self.bins = {} 100 | _parent_stack.append(self) 101 | try: 102 | for arg, (name, field) in zip(args, self.__class__._fields.items()): 103 | # if isinstance(field, BinarySchemaMetaclass): 104 | # binary = arg.binary 105 | # elif isinstance(field, Unit): 106 | # binary = field(arg) 107 | # if arg is ...: 108 | # arg = field.get_default() 109 | setattr(self, name, arg) 110 | # self.bins[name] = binary 111 | finally: 112 | _parent_stack.pop() 113 | 114 | if hasattr(self, "__post_init__"): 115 | self.__post_init__() 116 | 117 | def member_get(self, name): 118 | return self.values[name] 119 | 120 | def member_set(self, name, value, binary): 121 | self.bins[name] = binary 122 | self.values[name] = value 123 | self._modified = True 124 | 125 | @property 126 | def binary(self): 127 | if self._modified: 128 | self._binary = b"".join(self.bins.values()) 129 | self._modified = False 130 | return self._binary 131 | 132 | def __str__(self): 133 | sl = [] 134 | for name in self.__class__._fields: 135 | value = getattr(self, name) 136 | sl.append(f"{name}={value!r}") 137 | s = ", ".join(sl) 138 | return f"{self.__class__.__name__}({s})" 139 | 140 | def __repr__(self): 141 | return f"<{self}>" 142 | 143 | def __eq__(self, other) -> bool: 144 | if not isinstance(other, self.__class__): 145 | return False 146 | for name in self.__class__._fields: 147 | if getattr(self, name) != getattr(other, name): 148 | return False 149 | return True 150 | 151 | 152 | # FieldType = typing.Union[BinarySchemaMetaclass, Unit] 153 | FieldType = typing.Union[typing.Type[BinarySchema], Unit] 154 | 155 | 156 | class MemberDescriptor: 157 | __slots__ = ("key", "member") 158 | 159 | def __init__(self, key: str, member: FieldType): 160 | self.key = key 161 | self.member = member 162 | 163 | def __get__(self, obj: typing.Optional[BinarySchema], owner): 164 | if obj is None: 165 | return self.member 166 | return obj.member_get(self.key) 167 | 168 | def __set__(self, obj: BinarySchema, value): 169 | if isinstance(self.member, BinarySchemaMetaclass): 170 | binary = value.binary 171 | elif isinstance(self.member, Unit): 172 | binary = self.member(value) 173 | if value is ...: 174 | value = self.member.parse(binary) 175 | obj.member_set(self.key, value, binary) 176 | 177 | 178 | class StructUnit(Unit): 179 | def __init__(self, format_: str): 180 | self._struct = Struct(format_) 181 | 182 | def __str__(self): 183 | return f"{self.__class__.__name__}({self._struct.format})" 184 | 185 | def get_value(self): 186 | return (yield from read_raw_struct(self._struct))[0] 187 | 188 | def __call__(self, obj) -> bytes: 189 | return self._struct.pack(obj) 190 | 191 | 192 | class IntUnit(Unit): 193 | def __init__(self, length: int, byteorder: str, signed: bool = False): 194 | self.length = length 195 | self.byteorder = byteorder 196 | self.signed = signed 197 | 198 | def get_value(self): 199 | return ( 200 | yield from read_int( 201 | self.length, byteorder=self.byteorder, signed=self.signed 202 | ) 203 | ) 204 | 205 | def __call__(self, obj: int) -> bytes: 206 | return obj.to_bytes(self.length, self.byteorder, signed=self.signed) 207 | 208 | 209 | int8 = StructUnit("b") 210 | uint8 = StructUnit("B") 211 | int16 = StructUnit("h") 212 | int16be = StructUnit(">h") 213 | uint16 = StructUnit("H") 214 | uint16be = StructUnit(">H") 215 | int24 = IntUnit(3, "little", signed=True) 216 | int24be = IntUnit(3, "big", signed=True) 217 | uint24 = IntUnit(3, "little", signed=False) 218 | uint24be = IntUnit(3, "big", signed=False) 219 | int32 = StructUnit("i") 220 | int32be = StructUnit(">i") 221 | uint32 = StructUnit("I") 222 | uint32be = StructUnit(">I") 223 | int64 = StructUnit("q") 224 | int64be = StructUnit(">q") 225 | uint64 = StructUnit("Q") 226 | uint64be = StructUnit(">Q") 227 | float16 = StructUnit("e") 228 | float16be = StructUnit(">e") 229 | float32 = StructUnit("f") 230 | float32be = StructUnit(">f") 231 | float64 = StructUnit("d") 232 | float64be = StructUnit(">d") 233 | 234 | 235 | class Bytes(Unit): 236 | def __init__(self, length: int): 237 | self.length = length 238 | if length >= 0: 239 | self._struct = Struct(f"{length}s") 240 | 241 | def __str__(self): 242 | return f"{self.__class__.__name__}({self.length})" 243 | 244 | def get_value(self): 245 | if self.length >= 0: 246 | return (yield from read_raw_struct(self._struct))[0] 247 | else: 248 | return (yield from read()) 249 | 250 | def __call__(self, obj) -> bytes: 251 | if self.length >= 0: 252 | return self._struct.pack(obj) 253 | else: 254 | return obj 255 | 256 | 257 | class MustEqual(Unit): 258 | def __init__(self, unit: Unit, value: typing.Any): 259 | self.unit = unit 260 | self.value = value 261 | 262 | def __str__(self): 263 | return f"{self.__class__.__name__}({self.unit}, {self.value})" 264 | 265 | def get_value(self): 266 | result = yield from self.unit.get_value() 267 | if self.value != result: 268 | raise ValueError(f"expect {self.value}, got {result}") 269 | return result 270 | 271 | def __call__(self, obj) -> bytes: 272 | if obj is not ...: 273 | if self.value != obj: 274 | raise ValueError(f"expect {self.value}, got {obj}") 275 | return self.unit(self.value) 276 | 277 | 278 | class EndWith(Unit): 279 | def __init__(self, bytes_: bytes): 280 | self.bytes_ = bytes_ 281 | 282 | def __str__(self): 283 | return f"{self.__class__.__name__}({self.bytes_})" 284 | 285 | def get_value(self): 286 | return (yield from read_until(self.bytes_, return_tail=False)) 287 | 288 | def __call__(self, obj: bytes) -> bytes: 289 | return obj + self.bytes_ 290 | 291 | 292 | class LengthPrefixedBytes(Unit): 293 | def __init__(self, length_unit: typing.Union[StructUnit, IntUnit]): 294 | self.length_unit = length_unit 295 | 296 | def __str__(self): 297 | return f"{self.__class__.__name__}({self.length_unit})" 298 | 299 | def get_value(self): 300 | length = yield from self.length_unit.get_value() 301 | return (yield from read_struct(f"{length}s"))[0] 302 | 303 | def __call__(self, obj: bytes) -> bytes: 304 | length = len(obj) 305 | return self.length_unit(length) + struct.pack(f"{length}s", obj) 306 | 307 | 308 | class LengthPrefixed(Unit): 309 | def __init__( 310 | self, length_unit: typing.Union[StructUnit, IntUnit], object_unit: FieldType 311 | ): 312 | self.length_unit = length_unit 313 | self.object_unit = object_unit 314 | 315 | def __str__(self): 316 | return f"{self.__class__.__name__}({self.length_unit}, {self.object_unit})" 317 | 318 | def get_value(self): 319 | length = yield from self.length_unit.get_value() 320 | (data,) = yield from read_struct(f"{length}s") 321 | parser = Parser(self._gen()) 322 | return parser.parse(data) 323 | 324 | @abc.abstractmethod 325 | def _gen(self) -> typing.Generator: 326 | "" 327 | 328 | 329 | class LengthPrefixedObjectList(LengthPrefixed): 330 | def _gen(self): 331 | parser = yield from get_parser() 332 | lst = [] 333 | yield from wait() 334 | while parser.has_more_data(): 335 | lst.append((yield from self.object_unit.get_value())) 336 | return lst 337 | 338 | def __call__(self, obj_list: typing.List[FieldType]) -> bytes: 339 | if isinstance(self.object_unit, BinarySchemaMetaclass): 340 | bytes_ = b"".join(bs.binary for bs in obj_list) 341 | elif isinstance(self.object_unit, Unit): 342 | bytes_ = b"".join(self.object_unit(bs) for bs in obj_list) 343 | return self.length_unit(len(bytes_)) + bytes_ 344 | 345 | 346 | class LengthPrefixedObject(LengthPrefixed): 347 | def _gen(self): 348 | parser = yield from get_parser() 349 | v = yield from self.object_unit.get_value() 350 | if parser.has_more_data(): 351 | raise ValueError("extra bytes left") 352 | return v 353 | 354 | def __call__(self, obj: FieldType) -> bytes: 355 | bytes_ = ( 356 | obj.binary 357 | if isinstance(self.object_unit, BinarySchemaMetaclass) 358 | else self.object_unit(obj) 359 | ) 360 | return self.length_unit(len(bytes_)) + bytes_ 361 | 362 | 363 | class Switch(Unit): 364 | def __init__(self, ref: str, cases: typing.Mapping[typing.Any, FieldType]): 365 | self.ref = ref 366 | self.cases = cases 367 | 368 | def __str__(self): 369 | return f"{self.__class__.__name__}({self.ref}, {self.cases})" 370 | 371 | def get_value(self): 372 | parser = yield from get_parser() 373 | mapping = parser._mapping_stack[-1] 374 | unit = self.cases[mapping[self.ref]] 375 | return (yield from unit.get_value()) 376 | 377 | def __call__(self, obj) -> bytes: 378 | parent = _parent_stack[-1] 379 | real_field = self.cases[getattr(parent, self.ref)] 380 | return real_field(obj) if isinstance(real_field, Unit) else obj.binary 381 | 382 | 383 | class SizedIntEnum(Unit): 384 | def __init__( 385 | self, 386 | size_unit: typing.Union[StructUnit, IntUnit], 387 | enum_class: typing.Type[enum.IntEnum], 388 | ): 389 | self.size_unit = size_unit 390 | self.enum_class = enum_class 391 | 392 | def __str__(self): 393 | return f"{self.__class__.__name__}({self.size_unit}, {self.enum_class})" 394 | 395 | def get_value(self): 396 | v = yield from self.size_unit.get_value() 397 | return self.enum_class(v) 398 | 399 | def __call__(self, obj: enum.IntEnum) -> bytes: 400 | return self.size_unit(obj.value) 401 | 402 | 403 | class Convert(Unit): 404 | def __init__(self, unit: Unit, *, encode: typing.Callable, decode: typing.Callable): 405 | self.unit = unit 406 | self.encode = encode 407 | self.decode = decode 408 | 409 | def __str__(self): 410 | return ( 411 | f"{self.__class__.__name__}" 412 | f"({self.unit}, encode={self.encode}, decode={self.decode})" 413 | ) 414 | 415 | def get_value(self): 416 | v = yield from self.unit.get_value() 417 | return self.decode(v) 418 | 419 | def __call__(self, obj: typing.Any) -> bytes: 420 | return self.unit(self.encode(obj)) 421 | 422 | 423 | class String(Convert): 424 | def __init__(self, length: int, encoding="utf-8"): 425 | super().__init__( 426 | Bytes(length), 427 | encode=lambda x: x.encode(encoding), 428 | decode=lambda x: x.decode(encoding), 429 | ) 430 | 431 | 432 | class LengthPrefixedString(Convert): 433 | def __init__( 434 | self, length_unit: typing.Union[StructUnit, IntUnit], encoding="utf-8" 435 | ): 436 | super().__init__( 437 | LengthPrefixedBytes(length_unit), 438 | encode=lambda x: x.encode(encoding), 439 | decode=lambda x: x.decode(encoding), 440 | ) 441 | 442 | 443 | def Group(**fields: typing.Dict[str, FieldType]) -> typing.Type[BinarySchema]: 444 | return type("Group", (BinarySchema,), fields) 445 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.black] 2 | line-length = 88 3 | include = '\.pyi?$' 4 | 5 | [tool.isort] 6 | line_length = 88 7 | multi_line_output = 3 8 | include_trailing_comma = true 9 | 10 | [build-system] 11 | requires = ["flit_core >=2,<4"] 12 | build-backend = "flit_core.buildapi" 13 | 14 | [tool.flit.metadata] 15 | module = "iofree" 16 | author = "Yingbo Gu" 17 | author-email = "tensiongyb@gmail.com" 18 | maintainer="Yingbo Gu" 19 | maintainer-email="tensiongyb@gmail.com" 20 | home-page = "https://github.com/guyingbo/iofree" 21 | description-file = "README.md" 22 | requires-python=">=3.6" 23 | classifiers = [ 24 | "License :: OSI Approved :: MIT License", 25 | "Intended Audience :: Developers", 26 | "Programming Language :: Python :: 3.6", 27 | "Programming Language :: Python :: 3.7", 28 | "Programming Language :: Python :: 3.8", 29 | ] 30 | dev-requires = ["pytest", "coverage", "pytest-cov"] 31 | keywords = "parser Sans-IO io-free" 32 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | description-file = README.md 3 | 4 | [aliases] 5 | test = pytest 6 | 7 | [tool:pytest] 8 | addopts = --verbose tests --cov=iofree --cov=tests --cov-report=term-missing --cov-report=xml 9 | 10 | [isort] 11 | line_length = 88 12 | multi_line_output = 3 13 | include_trailing_comma = true 14 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_namespace_packages 2 | import os.path 3 | import re 4 | 5 | VERSION_RE = re.compile(r"""__version__ = ['"]([-a-z0-9.]+)['"]""") 6 | BASE_PATH = os.path.dirname(__file__) 7 | 8 | 9 | with open(os.path.join(BASE_PATH, "iofree", "__init__.py")) as f: 10 | try: 11 | version = VERSION_RE.search(f.read()).group(1) 12 | except IndexError: 13 | raise RuntimeError("Unable to determine version.") 14 | 15 | 16 | with open(os.path.join(BASE_PATH, "README.md")) as readme: 17 | long_description = readme.read() 18 | 19 | description = ( 20 | "An io-free stream parser " 21 | "which helps implementing network protocols in the `Sans-IO` way" 22 | ) 23 | 24 | setup( 25 | name="iofree", 26 | description=description, 27 | long_description=long_description, 28 | long_description_content_type="text/markdown", 29 | license="MIT", 30 | version=version, 31 | author="Yingbo Gu", 32 | author_email="tensiongyb@gmail.com", 33 | maintainer="Yingbo Gu", 34 | maintainer_email="tensiongyb@gmail.com", 35 | url="https://github.com/guyingbo/iofree", 36 | packages=find_namespace_packages(include=["iofree*"]), 37 | python_requires=">=3.6", 38 | classifiers=[ 39 | "License :: OSI Approved :: MIT License", 40 | "Framework :: AsyncIO", 41 | "Intended Audience :: Developers", 42 | "Programming Language :: Python :: 3.6", 43 | "Programming Language :: Python :: 3.7", 44 | "Programming Language :: Python :: 3.8", 45 | ], 46 | setup_requires=["pytest-runner"], 47 | tests_require=["pytest", "coverage", "pytest-cov"], 48 | ) 49 | -------------------------------------------------------------------------------- /tests/test_iofree.py: -------------------------------------------------------------------------------- 1 | import random 2 | from datetime import datetime 3 | 4 | import pytest 5 | 6 | import iofree 7 | from iofree import schema 8 | 9 | 10 | class HTTPResponse(schema.BinarySchema): 11 | head = schema.EndWith(b"\r\n\r\n") 12 | 13 | def __post_init__(self): 14 | first_line, *header_lines = self.head.split(b"\r\n") 15 | self.ver, self.code, *status = first_line.split(None, 2) 16 | self.status = status[0] if status else b"" 17 | self.header_lines = header_lines 18 | 19 | 20 | @iofree.parser 21 | def http_response(): 22 | response = yield from HTTPResponse 23 | assert response.ver == b"HTTP/1.1" 24 | assert response.code == b"200" 25 | assert response.status == b"OK" 26 | yield from iofree.read_until(b"\n", return_tail=True) 27 | data = yield from iofree.read(4) 28 | assert data == b"haha" 29 | (number,) = yield from iofree.read_struct("!H") 30 | assert number == 8 * 256 + 8 31 | number = yield from iofree.read_int(3) 32 | assert number == int.from_bytes(b"\x11\x11\x11", "big") 33 | assert (yield from iofree.peek(2)) == b"co" 34 | assert (yield from iofree.read(7)) == b"content" 35 | yield from iofree.wait() 36 | assert len((yield from iofree.read_more(5))) >= 5 37 | yield from iofree.read() 38 | yield from iofree.wait_event() 39 | return b"\r\n".join(response.header_lines) 40 | 41 | 42 | def test_http_parser(): 43 | parser = http_response.parser() 44 | response = bytearray( 45 | b"HTTP/1.1 200 OK\r\n" 46 | b"Connection: keep-alive\r\n" 47 | b"Content-Encoding: gzip\r\n" 48 | b"Content-Type: text/html\r\n" 49 | b"Date: " 50 | + datetime.now().strftime("%a, %d %b %Y %H:%M:%S GMT").encode() 51 | + b"\r\nServer: nginx\r\n" 52 | b"Vary: Accept-Encoding\r\n\r\n" 53 | b"a line\nhaha\x08\x08\x11\x11\x11content extra" 54 | ) 55 | while response: 56 | n = random.randrange(1, 30) 57 | data = response[:n] 58 | del response[:n] 59 | parser.send(data) 60 | parser.send() 61 | parser.send_event(0) 62 | assert parser.has_result 63 | parser.get_result() 64 | parser.read_output_bytes() 65 | parser.send(b"redundant") 66 | assert parser.readall().endswith(b"redundant") 67 | 68 | 69 | def test_http_parser2(): 70 | for i in range(100): 71 | test_http_parser() 72 | 73 | 74 | @iofree.parser 75 | def simple(): 76 | parser = yield from iofree.get_parser() 77 | yield from iofree.read(1) 78 | assert parser.has_more_data() 79 | assert not parser.finished() 80 | raise Exception("special") 81 | 82 | 83 | @iofree.parser 84 | def bad_reader(): 85 | with pytest.raises(ValueError): 86 | yield from iofree.read_more(-1) 87 | with pytest.raises(ValueError): 88 | yield from iofree.peek(-1) 89 | with pytest.raises(ValueError): 90 | yield from iofree.read_int(-1) 91 | 92 | yield from iofree.wait() 93 | yield "bad" 94 | 95 | 96 | def test_exception(): 97 | parser = simple.parser() 98 | with pytest.raises(iofree.ParseError): 99 | parser.send(b"haha") 100 | with pytest.raises(iofree.NoResult): 101 | parser.get_result() 102 | 103 | 104 | def test_bad(): 105 | parser = bad_reader.parser() 106 | with pytest.raises(RuntimeError): 107 | parser.send(b"haha") 108 | -------------------------------------------------------------------------------- /tests/test_parser.py: -------------------------------------------------------------------------------- 1 | import socket 2 | import threading 3 | from time import sleep 4 | 5 | import pytest 6 | 7 | import iofree 8 | from iofree import schema 9 | from iofree.contrib.common import Addr 10 | 11 | 12 | @iofree.parser 13 | def example(): 14 | parser = yield from iofree.get_parser() 15 | repr(parser) 16 | yield from Addr 17 | while True: 18 | n = yield from schema.uint8 19 | if n == 16: 20 | parser.respond(data=b"done", exc=Exception()) 21 | elif n == 32: 22 | parser.respond(close=True, result=10) 23 | 24 | 25 | def write_data(sock, n): 26 | sock.sendall(Addr.from_tuple(("google.com", 8080)).binary) 27 | for i in range(n): 28 | sock.sendall(i.to_bytes(1, "big")) 29 | sleep(0.001) 30 | sock.shutdown(socket.SHUT_WR) 31 | 32 | 33 | def test_parser(): 34 | parser = example.parser() 35 | rsock, wsock = socket.socketpair() 36 | thread = threading.Thread(target=write_data, args=(wsock, 0x21)) 37 | thread.start() 38 | with pytest.raises(Exception): 39 | r = parser.run(rsock) 40 | r = parser.run(rsock) 41 | assert r == 10 42 | thread.join() 43 | 44 | 45 | def test_parser2(): 46 | parser = example.parser() 47 | rsock, wsock = socket.socketpair() 48 | thread = threading.Thread(target=write_data, args=(wsock, 5)) 49 | thread.start() 50 | with pytest.raises(iofree.ParseError): 51 | parser.run(rsock) 52 | thread.join() 53 | 54 | 55 | @iofree.parser 56 | def first(): 57 | parser = yield from iofree.get_parser() 58 | for i in range(10): 59 | a = yield from schema.Group(x=schema.uint8, y=schema.uint16be) 60 | parser.respond(result=a.binary[1:] + a.binary[:1]) 61 | return b"" 62 | 63 | 64 | @iofree.parser 65 | def second(): 66 | parser = yield from iofree.get_parser() 67 | for i in range(10): 68 | a = yield from schema.Group(x=schema.uint16be, y=schema.uint8) 69 | parser.respond(result=a) 70 | 71 | 72 | def test_parser_chain(): 73 | p = iofree.ParserChain(first.parser(), second.parser()) 74 | p.send(b"") 75 | for i in range(10): 76 | p.send(schema.Group(x=schema.uint8, y=schema.uint16be)(30, 512).binary) 77 | list(p) 78 | -------------------------------------------------------------------------------- /tests/test_schema.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from iofree import schema 4 | 5 | 6 | def check_schema(schema): 7 | str(schema) 8 | repr(schema) 9 | str(schema.__class__) 10 | schema2 = schema.__class__.parse(schema.binary) 11 | assert schema == schema2 12 | 13 | 14 | def test_number(): 15 | assert schema.uint16be.parse(b"\x00\x03") == 3 16 | with pytest.raises(schema.ParseError): 17 | schema.uint8.parse(b"\x03\x04", strict=True) 18 | 19 | 20 | def test_length_prefixed_bytes(): 21 | some_bytes = schema.LengthPrefixedBytes(schema.uint8) 22 | str(some_bytes) 23 | 24 | 25 | def test_basic(): 26 | class Content(schema.BinarySchema): 27 | first_line = schema.EndWith(b"\r\n") 28 | string = schema.String(5, encoding="ascii") 29 | 30 | content = Content(b"GET / HTTP/1.1", "abcde") 31 | assert content != 3 32 | content2 = Content(b"POST", "xyz") 33 | assert content != content2 34 | with pytest.raises(ValueError): 35 | Content(b"PUT") 36 | check_schema(content) 37 | 38 | with pytest.raises(schema.ParseError): 39 | Content.parse(b"abc\r\n" + "中文".encode()) 40 | 41 | 42 | def test_equal(): 43 | class Content(schema.BinarySchema): 44 | name = schema.MustEqual(schema.Bytes(3), b"abc") 45 | name2 = schema.LengthPrefixedObject(schema.uint8, schema.Bytes(1)) 46 | 47 | with pytest.raises(schema.ParseError): 48 | Content.parse(b"abb") 49 | 50 | with pytest.raises(schema.ParseError): 51 | Content.parse(b"abc\x02abc") 52 | 53 | 54 | def test_group(): 55 | class Dynamic(schema.BinarySchema): 56 | a = schema.uint8 57 | b = schema.Group(c=schema.uint16, d=schema.uint24be) 58 | 59 | dynamic = Dynamic(1, Dynamic.b(2, 3)) 60 | check_schema(dynamic) 61 | 62 | 63 | def test_schema(): 64 | G = schema.Group(a=schema.uint8, b=schema.int32) 65 | Dynamic = schema.Group( 66 | a=schema.LengthPrefixedBytes(schema.uint24be), 67 | b=schema.LengthPrefixedObject(schema.uint32, schema.EndWith(b"\n")), 68 | c=schema.LengthPrefixedObjectList(schema.uint16, schema.String(3)), 69 | d=schema.LengthPrefixedObjectList(schema.uint64, G), 70 | ) 71 | dynamic = Dynamic(b"abc", b"def", ["123", "456"], [G(3, 5), G(6, 10)]) 72 | check_schema(dynamic) 73 | -------------------------------------------------------------------------------- /tests/test_socks5.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import pytest 4 | 5 | from iofree.contrib import socks5 6 | 7 | 8 | def check_schema(schema): 9 | str(schema) 10 | repr(schema) 11 | str(schema.__class__) 12 | parser = schema.__class__.get_parser() 13 | schema2 = parser.parse(schema.binary) 14 | assert schema == schema2 15 | 16 | 17 | def test_handshake(): 18 | handshake = socks5.Handshake( 19 | ..., [socks5.AuthMethod.no_auth, socks5.AuthMethod.user_auth] 20 | ) 21 | check_schema(handshake) 22 | with pytest.raises(ValueError): 23 | socks5.Handshake(6, []) 24 | 25 | 26 | def test_client_request(): 27 | auth = socks5.UsernameAuth(..., "username", "password") 28 | check_schema(auth) 29 | 30 | request = socks5.ClientRequest( 31 | ..., socks5.Cmd.connect, 0, socks5.Addr(1, "127.0.0.1", 8000) 32 | ) 33 | check_schema(request) 34 | 35 | 36 | def test_reply(): 37 | reply = socks5.Reply(5, socks5.Rep.succeeded, 0, socks5.Addr(4, "::1", 8080)) 38 | check_schema(reply) 39 | 40 | 41 | def test_udp_reply(): 42 | udp_reply = socks5.UDPRelay( 43 | ..., 32, socks5.Addr.from_tuple(("google.com", 80)), os.urandom(64) 44 | ) 45 | check_schema(udp_reply) 46 | --------------------------------------------------------------------------------