# Source code for nobodd.tools

    # nobodd: a boot configuration tool for the Raspberry Pi
    #
    # Copyright (c) 2023-2026 Dave Jones <dave.jones@canonical.com>
    # Copyright (c) 2023-2026 Canonical Ltd.
    #
    # SPDX-License-Identifier: GPL-3.0

    import io
    import codecs
    import socket
    import datetime as dt
    from collections.abc import Mapping
    from contextlib import contextmanager

    from . import lang

    def labels(desc):
        """
        Given the description of a C structure in *desc*, returns a tuple of the
        labels.

        The :class:`str` *desc* must contain one entry per line (blank lines are
        ignored) where each entry consists of whitespace separated type (in Python
        :mod:`struct` format) and label. For example::

            >>> EBPB = '''
            B     drive_number
            1x    reserved
            B     extended_boot_sig
            4s    volume_id
            11s   volume_label
            8s    file_system
            '''
            >>> labels(EBPB)
            ('drive_number', 'extended_boot_sig', 'volume_id', 'volume_label',
            'file_system')

        Note the amount of whitespace is arbitrary, and further that any entries
        with the type "x" (which is used to indicate padding) will be excluded from
        the result ("reserved" is missing from the result tuple above).

        The corresponding function :func:`formats` can be used to obtain a tuple
        of the types.
        """
        return tuple(
            label
            for line in desc.splitlines()
            if line
            for fmt, label in (line.split(None, 1),)
            if not fmt.endswith('x')
        )

    def formats(desc, prefix='<'):
        """
        Given the description of a C structure in *desc*, returns a concatenated
        :class:`str` of the types with an optional *prefix* (for endianness).

        The :class:`str` *desc* must contain one entry per line (blank lines are
        ignored) where each entry consists of whitespace separated type (in Python
        :mod:`struct` format) and label. For example::

            >>> EBPB = '''
            B     drive_number
            1x    reserved
            B     extended_boot_sig
            4s    volume_id
            11s   volume_label
            8s    file_system
            '''
            >>> formats(EBPB)
            '<B1xB4s11s8s'

        Note the amount of whitespace is arbitrary, and further that any entries
        with the type "x" (which is used to indicate padding) are *not* excluded
        (unlike in :func:`labels`).

        The corresponding function :func:`labels` can be used to obtain a tuple
        of the labels.
        """
        return prefix + ''.join(
            fmt
            for line in desc.splitlines()
            if line
            for fmt, label in (line.split(None, 1),)
        )

    def get_best_family(host, port):
        """
        Given a *host* name and a *port* specification (either a number or a
        service name), returns the network family (e.g. ``socket.AF_INET``) and
        socket address to listen on as a tuple.
        """
        try:
            infos = socket.getaddrinfo(
                host, port, proto=socket.IPPROTO_UDP, flags=socket.AI_PASSIVE)
        except socket.gaierror as exc:
            raise ValueError(lang._('invalid host and port combination')) from exc
        for family, _, _, _, sockaddr in infos:
            return family, sockaddr
        raise ValueError(lang._('invalid host and port combination'))

    def format_address(address):
        """
        Given a socket *address*, return a suitable :class:`str` representation of
        it.

        Specifically, for IP4 addresses a simple "host:port" representation is
        used. For IP6 addresses (which typically incorporate ":" in the host
        portion), a "[host]:port" variant is used.
        """
        host, port, *extra = address
        if ':' in host:
            return f'[{host}]:{port}'
        else:
            return f'{host}:{port}'

    class BufferedTranscoder(io.RawIOBase):
        """
        A read-only transcoder, somewhat similar to :class:`codecs.StreamRecoder`,
        but which strictly obeys the definition of the ``read`` method (with
        internal buffering).

        This class is primarily intended for use in :mod:`~nobodd.netascii` encoded
        transfers where it is used to transcode the underlying file stream into
        netascii encoding for the TFTP server.

        The built-in :class:`codecs.StreamRecoder` class would seem ideal for this
        but for one issue: under certain circumstances (including those involved in
        netascii encoding), it violates the contract of the ``read`` method by
        returning *more* bytes than requested. For example::

            >>> import io, codecs
            >>> latin1_stream = io.BytesIO('abcdé'.encode('latin-1'))
            >>> utf8_stream = codecs.StreamRecoder(latin1_stream,
            ... codecs.getencoder('utf-8'), codecs.getdecoder('utf-8'),
            ... codecs.getreader('latin-1'), codecs.getwriter('latin-1'))
            >>> utf8_stream.read(3)
            b'abc'
            >>> utf8_stream.read(1)
            b'd'
            >>> utf8_stream.read(1)
            b'\\xc3\\xa9'

        This is alluded to in the documentation of :class:`StreamReader.read` so it
        probably isn't a bug, but it is rather inconvenient when the caller is
        looking to fill a network packet of a specific size, and thus expects not
        to over-run.

        This class implements a rather simpler recoder, which is read-only, does
        not permit seeking, but by use of an internal buffer, guarantees that the
        :meth:`read` method (and associated methods like :meth:`readinto`) will
        not return more bytes than requested.

        It is constructed with the underlying *stream*, the name of the
        *output_encoding*, the name of the *input_encoding* (which defaults to the
        *output_encoding* when not specified), and the *errors* mode to use with
        the codecs. For example::

            >>> import io
            >>> from nobodd.tools import BufferedTranscoder
            >>> latin1_stream = io.BytesIO('abcdé'.encode('latin-1'))
            >>> utf8_stream = BufferedTranscoder(latin1_stream, 'utf-8', 'latin-1')
            >>> utf8_stream.read(4)
            b'abcd'
            >>> utf8_stream.read(1)
            b'\\xc3'
            >>> utf8_stream.read(1)
            b'\\xa9'
        """
        def __init__(self, stream, output_encoding, input_encoding=None,
                     errors='strict'):
            if input_encoding is None:
                input_encoding = output_encoding
            self._source = codecs.getreader(input_encoding)(stream, errors)
            self._encode = codecs.getencoder(output_encoding)
            self._buffer = bytearray()

        def readable(self):
            return True

        def readall(self):
            result = self._buffer + self._encode(self._source.read())[0]
            del self._buffer[:]
            return result

        def readinto(self, b):
            while len(self._buffer) < len(b):
                s = self._source.read(4096)
                if not s:
                    break
                self._buffer.extend(self._encode(s)[0])
            to_read = min(len(b), len(self._buffer))
            b[:to_read] = self._buffer[:to_read]
            del self._buffer[:to_read]
            return to_read

    class FrozenDict(Mapping):
        """
        A hashable, immutable mapping type.

        The arguments to :class:`FrozenDict` are processed just like those to
        :class:`dict`.
        """
        def __init__(self, *args):
            self._d = dict(*args)
            self._hash = None

        def __iter__(self):
            return iter(self._d)

        def __len__(self):
            return len(self._d)

        def __getitem__(self, key):
            return self._d[key]

        def __repr__(self):
            return f'{self.__class__.__name__}({self._d})'

        def __hash__(self):
            if self._hash is None:
                self._hash = hash((frozenset(self), frozenset(self.values())))
            return self._hash

    def decode_timestamp(date, time, cs=0, tz=dt.timezone.utc):
        """
        Given the integers *date*,  *time*, and optionally *cs* (from various
        fields in :class:`~nobodd.fat.DirectoryEntry`), return a
        :class:`~datetime.datetime` with the decoded timestamp. The returned
        value will have the :class:`~datetime.tzinfo` set by *tz*.
        """
        ms = cs * 10
        return dt.datetime(
            year=1980 + ((date & 0xFE00) >> 9),
            month=(date & 0x1E0) >> 5,
            day=(date & 0x1F),
            hour=(time & 0xF800) >> 11,
            minute=(time & 0x7E0) >> 5,
            second=(time & 0x1F) * 2 + (ms // 1000),
            microsecond=(ms % 1000) * 1000,
            tzinfo=tz
        )

    def encode_timestamp(ts):
        """
        Given a :class:`~datetime.datetime`, encode it as a FAT-compatible triple
        of three 16-bit integers representing (date, time, 1/100th seconds).
        """
        earliest = dt.datetime(1980, 1, 1, tzinfo=ts.tzinfo)
        latest = dt.datetime(2100, 1, 1, tzinfo=ts.tzinfo)
        if not earliest <= ts < latest:
            raise ValueError(f'{ts} is outside the valid range for FAT timestamps')
        return (
            ((ts.year - 1980) << 9) | (ts.month << 5) | ts.day,
            (ts.hour << 11) | (ts.minute << 5) | (ts.second // 2),
            ((ts.second % 2) * 1000 + (ts.microsecond // 1000)) // 10
        )

    def exclude(ranges, value):
        """
        Given a list non-overlapping of *ranges*, sorted in ascending order, this
        function modifies the range containing *value* (an integer, which must
        belong to one and only one range in the list) to exclude it.
        """
        for i, r in enumerate(ranges):
            if value in r:
                break
        else:
            return
        ranges[i:i + 1] = [
            r for r in (range(r.start, value), range(value + 1, r.stop)) if r]

    def any_match(s, expressions):
        """
        Given a :class:`str` *s*, and *expressions*, a sequence of compiled
        regexes, return the :class:`re.Match` object from the first regex that
        matches *s*. If no regexes match, return :data:`None`.
        """
        for exp in expressions:
            m = exp.match(s)
            if m:
                return m
        return None

    @contextmanager
    def open_file(arg, mode='w'):
        """
        A context manager for handling files specified on the command line. Handles
        returning a standard stream when *arg* is "-", and closes files when the
        context exits (if required).
        """
        if arg is None:
            file = None
        elif arg == '-':
            file = sys.stdin if set(mode) & set('r+') == {'r'} else sys.stdout
            file = file.buffer if 'b' in mode else file
        else:
            file = open(arg, mode)
        try:
            yield file
        finally:
            if arg is not None and arg != '-':
                file.close()