asn1 app over soup (#44)

Author: SamDanielThangarajan

.pylintrc

@@ -16,4 +16,5 @@ disable=
     W0707, #too-many-statements
     R0801, #similarities
     R0913, #too-many-arguments
-    R0917, #too-many-positional-arguments
+    R0917, #too-many-positional-arguments
+    W1203, #logging-fstring-interpolation

pyproject.toml

@@ -16,7 +16,9 @@ classifiers = [
 dependencies = [
     'attrs>=23.1',
     'chevron>=0.14.0',
-    'click>=8.1'
+    'click>=8.1',
+    'asn1>=2.6',
+    'asn1tools==0.164.0'
 ]
 dynamic = ["version"]
 

requirements.txt

@@ -3,4 +3,6 @@ chevron>=0.14.0
 click>=8.1
 pytest>=8.3
 pytest-cov>=4
+asn1>=2.6
+asn1tools==0.164.0
 pytest-asyncio>=0.23

src/nasdaq_protocols/asn1_app/__init__.py

@@ -0,0 +1,55 @@
+from typing import Callable
+from nasdaq_protocols import soup
+
+from .soup_session import (
+    OnAns1CloseCoro,
+    OnAsn1MessageCoro,
+    Ans1SoupSessionId,
+    Asn1SoupClientSession
+)
+from .core import (
+    Asn1Message,
+    Asn1Spec
+)
+from .codegen import generate_soup_app
+
+
+__all__ = [
+    'Asn1Message',
+    'Asn1Spec',
+    'OnAns1CloseCoro',
+    'OnAsn1MessageCoro',
+    'Ans1SoupSessionId',
+    'generate_soup_app',
+    'connect_async_soup'
+]
+
+
+async def connect_async_soup(remote: tuple[str, int], user: str, passwd: str, session_id,  # pylint: disable=R0913
+                             session_factory: Callable[[soup.SoupClientSession], Asn1SoupClientSession],
+                             sequence: int = 0,
+                             client_heartbeat_interval: int = 10,
+                             server_heartbeat_interval: int = 10):
+    """
+    Connect to the Soup-Asn1 server.
+
+    :param remote: tuple of host and port
+    :param user: Username to login
+    :param passwd:  Password to login
+    :param session_id: Name of the session to join [Default=''] .
+    :param session_factory: Factory to create a ClientSession.
+    :param sequence: The sequence number. [Default=1]
+    :param client_heartbeat_interval: seconds between client heartbeats.
+    :param server_heartbeat_interval: seconds between server heartbeats.
+    :return: ClientSession
+    """
+
+    # Create a soup session
+    soup_session = await soup.connect_async(
+        remote, user, passwd, session_id, sequence=sequence,
+        client_heartbeat_interval=client_heartbeat_interval,
+        server_heartbeat_interval=server_heartbeat_interval
+    )
+
+    # Create an asn1-soup-session (with the soup session created above)
+    return session_factory(soup_session)

src/nasdaq_protocols/asn1_app/codegen.py

@@ -0,0 +1,114 @@
+import os
+import logging
+import shutil
+from importlib import resources
+from pathlib import Path
+
+import attrs
+import chevron
+import click
+
+from . import templates
+
+
+__all__ = [
+    'generate_soup_app'
+]
+TEMPLATES_PATH = resources.files(templates)
+LOG = logging.getLogger(__name__)
+
+
+@attrs.define(auto_attribs=True)
+class Ans1Generator:
+    asn1_input_files_dir: str
+    app_name: str
+    pdu: str
+    op_dir: str
+    template: str
+    prefix: str = ''
+    package_name: str = ''
+    generate_init_file: bool = False
+    _op_file: str = None
+    _module_name: str = None
+    _asn1_files: list = None
+
+    def __attrs_post_init__(self):
+        prefix = f'{self.prefix}_' if self.prefix else ''
+        self.template = self.template if self.template.endswith('.mustache') else f'{self.template}.mustache'
+        self.template = os.path.join(str(TEMPLATES_PATH), self.template)
+        self._module_name = f'{prefix}{self.app_name}'
+        self._op_file = os.path.join(self.op_dir, f'{self._module_name}.py')
+        self._asn1_files = [
+            os.path.join(self.asn1_input_files_dir, file)
+            for file in os.listdir(self.asn1_input_files_dir)
+            if file.endswith('.asn1')
+        ]
+        shutil.rmtree(self.op_dir)
+        Path(self.op_dir).mkdir(parents=True, exist_ok=True)
+
+    def generate(self, extra_context=None):
+        context = {
+            'app_name': self.app_name,
+            'pdu_name': self.pdu,
+            'package_name': self.package_name,
+            'spec': {
+                'name': self.app_name,
+                'capitalised_name': self.app_name.capitalize()
+            },
+        }
+        context.update(extra_context or {})
+
+        generated_files = [Ans1Generator._generate(self.template, context, self._op_file)]
+        if self.generate_init_file:
+            generated_files.append(self._generate_init_file())
+
+        self._copy_asn1_files()
+
+        return generated_files
+
+    def _generate_init_file(self):
+        init_file = os.path.join(self.op_dir, '__init__.py')
+        context = {
+            'modules': [
+                {'name': self._module_name}
+            ]
+        }
+        init_template = os.path.join(str(TEMPLATES_PATH), 'init.mustache')
+        Ans1Generator._generate(init_template, context, init_file)
+        return init_file
+
+    def _copy_asn1_files(self):
+        op_spec_dir = os.path.join(self.op_dir, 'spec')
+        Path(op_spec_dir).mkdir(parents=True, exist_ok=True)
+        for spec_file in self._asn1_files:
+            shutil.copy2(spec_file, os.path.join(op_spec_dir, os.path.basename(spec_file)))
+
+    @staticmethod
+    def _generate(template, context, op_file):
+        with open(op_file, 'a', encoding='utf-8') as op, open(template, 'r', encoding='utf-8') as inp:
+            code_as_string = chevron.render(inp.read(), context, partials_path=str(TEMPLATES_PATH))
+            op.write(code_as_string)
+            print(f'Generated: {op_file}')
+            return op_file
+
+
+@click.command()
+@click.option('--asn1-files-dir', type=click.Path(exists=True, file_okay=False, dir_okay=True, readable=True))
+@click.option('--app-name', type=click.STRING)
+@click.option('--pdu-name', type=click.STRING)
+@click.option('--prefix', type=click.STRING, default='')
+@click.option('--op-dir', type=click.Path(exists=True, writable=True))
+@click.option('--package-name', type=click.STRING)
+@click.option('--init-file/--no-init-file', show_default=True, default=True)
+def generate_soup_app(asn1_files_dir, app_name, pdu_name, prefix, op_dir, package_name, init_file):
+    generator = Ans1Generator(
+        asn1_files_dir,
+        app_name,
+        pdu_name,
+        op_dir,
+        template='ans1_soup_app.mustache',
+        prefix=prefix,
+        package_name=package_name,
+        generate_init_file=init_file
+    )
+    generator.generate()

src/nasdaq_protocols/asn1_app/core.py

@@ -0,0 +1,79 @@
+import importlib
+import logging
+from typing import ClassVar
+
+import asn1tools
+import attrs
+
+from asn1tools import CompileError, DecodeError, ParseError
+from nasdaq_protocols.common import Serializable, DuplicateMessageException
+
+__all__ = [
+    'Asn1Spec',
+    'Asn1Message'
+]
+LOG = logging.getLogger(__name__)
+
+
+class Asn1Spec:
+    SpecMap = {}
+
+    SpecName: str
+    SpecPkgDir: str
+    Spec: any
+
+    def __init_subclass__(cls, **kwargs):
+        try:
+            cls.SpecName, cls.SpecPkgDir = kwargs['spec_name'], kwargs['spec_pkg_dir']
+        except KeyError:
+            raise AttributeError('Missing spec_name or spec_pkg_dir')
+
+        if cls.SpecName in Asn1Spec.SpecMap:
+            spec_name = cls.SpecName
+            LOG.error(f'asn1 spec {spec_name} exists, existing spec = {Asn1Spec.SpecMap[spec_name]}')
+            raise DuplicateMessageException(existing_msg=Asn1Spec.SpecMap[spec_name], new_msg=cls)
+
+        try:
+            cls.Spec = Asn1Spec._compile(cls.SpecPkgDir)
+            Asn1Spec.SpecMap[cls.SpecName] = cls
+        except (CompileError, ParseError) as error:
+            LOG.error(f'Compile error - {cls.SpecName}, {error=}')
+            raise error
+
+    @classmethod
+    def decode(cls, pdu_name: str, bytes_: bytes) -> tuple[int, dict[any, any]]:
+        try:
+            decoded, length = cls.Spec.decode_with_length(pdu_name, bytes_)
+            return length, decoded
+        except DecodeError as err:
+            LOG.error(f'error while decoding asn1- {err=}')
+            return 0, {}
+
+    @staticmethod
+    def _compile(spec_pkg_dir):
+        files = list(
+            filter(
+                lambda x: x.suffix == '.asn1',
+                importlib.resources.files(spec_pkg_dir).iterdir()
+            )
+        )
+        return asn1tools.compiler.compile_files(files)
+
+
+@attrs.define
+class Asn1Message(Serializable):
+    PduName: ClassVar[str]
+    Spec: ClassVar[Asn1Spec]
+
+    def __init_subclass__(cls, **kwargs):
+        if all(_ in kwargs for _ in ['spec', 'pdu_name']):
+            cls.PduName = kwargs['pdu_name']
+            cls.Spec = kwargs['spec']
+            LOG.info(f'subclassed Asn1Message, {cls.Spec=}, {cls.PduName=}')
+
+    def to_bytes(self):
+        raise NotImplementedError('Encoding asn1 message is not supported yet.')
+
+    @classmethod
+    def from_bytes(cls, bytes_: bytes):
+        return cls.Spec.decode(cls.PduName, bytes_)

src/nasdaq_protocols/asn1_app/soup_session.py

@@ -0,0 +1,98 @@
+import asyncio
+from typing import Callable, Type, Awaitable, ClassVar
+
+import attrs
+from nasdaq_protocols.common import DispatchableMessageQueue, logable
+from nasdaq_protocols import soup
+
+from .core import Asn1Message
+
+
+__all__ = [
+    'OnAsn1MessageCoro',
+    'OnAns1CloseCoro',
+    'Ans1SoupSessionId',
+    'Asn1SoupClientSession'
+]
+OnAsn1MessageCoro = Callable[[Type[Asn1Message]], Awaitable[None]]
+OnAns1CloseCoro = Callable[[], Awaitable[None]]
+
+
+@attrs.define(auto_attribs=True)
+class Ans1SoupSessionId:
+    soup_session_id: soup.SoupSessionId = None
+
+    def __str__(self):
+        if self.soup_session_id:
+            return f'asn1-{self.soup_session_id}'
+        return 'asn1-nosoup'
+
+
+@attrs.define(auto_attribs=True)
+@logable
+class Asn1SoupClientSession:
+    """ Soup client session that can read messages from soup server.
+
+    Currently, only one pdu is supported.
+    """
+    Asn1Message: ClassVar[Asn1Message]
+    soup_session: soup.SoupClientSession
+    on_msg_coro: OnAsn1MessageCoro = None
+    on_close_coro: OnAns1CloseCoro = None
+    closed: bool = False
+    _session_id: Ans1SoupSessionId = None
+    _close_event: asyncio.Event = None
+    _message_queue: DispatchableMessageQueue = None
+
+    def __init_subclass__(cls, **kwargs):
+        if 'asn1_message' not in kwargs:
+            raise AttributeError("Missing 'asn1_message' attribute")
+        cls.Asn1Message = kwargs['asn1_message']
+
+    def __attrs_post_init__(self):
+        self._session_id = Ans1SoupSessionId(self.soup_session.session_id)
+        self._message_queue = DispatchableMessageQueue(self._session_id, self.on_msg_coro)
+        self.soup_session.set_handlers(on_msg_coro=self._on_soup_message, on_close_coro=self._on_soup_close)
+        self.soup_session.start_dispatching()
+
+    async def receive_message(self):
+        """
+        Asynchronously receive a message from the soup session.
+
+        This method blocks until a message is received by the session.
+        """
+        return await self._message_queue.get()
+
+    async def close(self):
+        """
+        Asynchronously close the session.
+        """
+        if self._close_event or self.closed:
+            self.log.debug('%s> closing in progress..', self._session_id)
+            return
+        self._close_event = asyncio.Event()
+        self.soup_session.initiate_close()
+        await self._close_event.wait()
+        self.log.debug('%s> closed.', self._session_id)
+
+    async def _on_soup_message(self, message: soup.SoupMessage):
+        if isinstance(message, soup.SequencedData):
+            self.log.debug('%s> incoming sequenced bytes_', self._session_id)
+            await self._message_queue.put(
+                self.decode(message.data)[1]
+            )
+
+    async def _on_soup_close(self):
+        await self._message_queue.stop()
+        if self.on_close_coro is not None:
+            await self.on_close_coro()
+        if self._close_event:
+            self._close_event.set()
+        self.closed = True
+
+    @classmethod
+    def decode(cls, bytes_: bytes):
+        """
+        Decode the given bytes into an asn1 message.
+        """
+        return cls.Asn1Message.from_bytes(bytes_)