<commit_before>import array
import contextlib
import ctypes
import fcntl
import io
import logging
import mmap
import os
import signal
import socket
import socketserver
import struct
import sys
import threading
from io import FileIO

from collections.abc import Container
from typing import Dict, Generator, List, Optional, Sequence, Tuple, TypeVar

from hades.bin.dhcp_script import main
from hades.common.signals import install_handler

logger = logging.getLogger(__name__)
SIZEOF_INT = ctypes.sizeof(ctypes.c_int)
memfd_create = ctypes.cdll.LoadLibrary("libc.so.6").memfd_create
MFD_CLOEXEC = 1
T = TypeVar('T')
Parser = Generator[int, Tuple[mmap.mmap, int], T]

class BaseParseError(Exception):
    def __init__(
            self,
            *args,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.element = element
        self.offset = offset
        super().__init__(*args, element, offset)

    def _prefix(self) -> str:
        offset, element = self.offset, self.element
        return "".join([
            "" if offset is None else "at offset {}: ".format(offset),
            "" if element is None else "while parsing {}: ".format(element),
        ]).capitalize()

    def with_element(self, element: str):
        self.element = element
        return self

    def with_offset(self, offset: int):
        self.offset = offset
        return self

class ParseError(BaseParseError):
    def __init__(
            self,
            message: str,
            *,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.element = element
        self.offset = offset
        self.message = message
        super().__init__(message, element=element, offset=offset)

    def __str__(self) -> str:
        return self._prefix() + self.message

class UnexpectedEOFError(BaseParseError):
    def __init__(
            self,
            needed: int,
            available: int,
            *,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.needed = needed
        self.available = available
        super().__init__(
            needed,
            available,
            element=element,
            offset=offset,
        )

    def __str__(self) -> str:
        return (
            "{}Unexpected end of file, expected at least {} more byte(s), "
            "but only {} byte(s) left."
            .format(self._prefix(), self.needed, self.available)
        )

class BufferTooSmallError(BaseParseError):
    def __init__(
            self,
            needed: int,
            available: int,
            *,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.needed = needed
        self.available = available
        super().__init__(
            needed,
            available,
            element=element,
            offset=offset,
        )

    def __str__(self) -> str:
        return (
            "{}Parser requires more data ({}) than can be buffered ({})."
            .format(self._prefix(), self.needed, self.available)
        )

class ProtocolError(Exception):
    pass

class Server(socketserver.UnixStreamServer):
    """
    Process :program:`dnsmasq` :option:`--dhcp-script` invocations.

    :program:`dnsmasq` can notify external tools about DHCP lease activity and
    even maintain a completely external lease database, if additionally
    :option:`--leasefile-ro` is specified. We use this mechanism to store the
    leases in the database.

    Starting a Python script that connects to a PostgreSQL database for DHCP
    lease activity of dnsmasq is too slow however:

    * The Python interpreter has some startup latency, but more importantly
      setuptools console scripts have a very long startup time (>1s) due to
      `this issue <https://github.com/pypa/setuptools/issues/510>`_.
    * PostgreSQL forks for every connection.

    To alleviate these issues we use a small and quick C client that passes its
    command-line arguments, environment variables and file descriptors over a
    UNIX socket to a long running Python server application that's permanently
    connected to the database. The server is single threaded and will only
    handle a single request at a time, because :program:`dnsmasq` itself is
    single threaded.

    The protocol is as follows:

    * We use a SOCK_STREAM AF_UNIX socket, so that the client knows on the
      socket level, when server has received all data and finished handling the
      request.
    * At first the client sends all data to server. The data may not exceed
      :py:code:`mmap.PAGESIZE - 1`. The content of the data is as follows:
        1. The first bytes are the number of arguments (:c:data:`argc`) as
           native type :c:type:`int`.
        2. This is followed by the contents of :c:data:`argv` as a series of
           zero terminated strings.
        3. After that the number of environment variables, that start with the
           prefix :envvar:`DNSMASQ_*` as native type :c:type:`int`.
        4. Followed by the actual environment variables as a series of zero
           terminated strings.
    * The three standard file descriptors stdin, stdout, and stderr of the
      script should be passed with data at some point in time via a
      :c:data:`SCM_RIGHTS` control message.
    * After all data has been sent, the client shuts down its write end of the
      connection and signals the server thereby, that it can begin processing
      the message.
    * The server will process the message and if necessary read additional data
      from the passed stdin file descriptor and write data to the passed stdout
      and stderr file descriptors.
    * After it has processed the message, the server sends a single byte status
      code between 0 and 127 and will close the connection. The script will
      exit with the status code.
    """
    max_packet_size = mmap.PAGESIZE - 1

    def __init__(self, sock, engine):
        self.engine = engine
        server_address = sock.getsockname()
        super().__init__(
            server_address, self._request_handler, bind_and_activate=False,
        )
        self.socket = sock
        # Leave one byte extra for trailing zero byte
        # TODO: With Python 3.8 a memfd can be opened and mapped twice:
        # writable and readonly
        fd = memfd_create(b"buffer", MFD_CLOEXEC)
        os.ftruncate(fd, self.max_packet_size + 1)
        self.buffer = mmap.mmap(
            fd,
            self.max_packet_size + 1,
            mmap.MAP_PRIVATE,
            mmap.PROT_READ | mmap.PROT_WRITE,
        )

    def _request_handler(
            self,
            request: socket.socket,
            client_address,
            server: socketserver.BaseServer,
    ):
        assert self == server

        logger.debug("Received new request from %s", client_address)
        (stdin, stdout, stderr), args, env = self._receive(request)
        status = os.EX_SOFTWARE
        try:
            status = self._process(stdin, stdout, stderr, args, env)
        finally:
            stdout.flush()
            stderr.flush()
            request.send(status.to_bytes(1, sys.byteorder))

    def _receive(
            self,
            request: socket.socket,
    ) -> Tuple[Tuple[FileIO, FileIO, FileIO], List[bytes], Dict[bytes, bytes]]:
        streams: List[FileIO] = []
        # Offset of the buffer relative to the input stream
        offset = 0
        # Number of filled bytes in the buffer
        available = 0
        # Initialize parser
        buffer = self.buffer
        buffer.seek(0, os.SEEK_SET)
        parser = self.parse_request(buffer, 0)
        needed = next(parser)
        with contextlib.ExitStack() as stack:
            while needed:
                # Prepare buffer for refill
                parsed = buffer.tell()
                offset += parsed
                buffer.move(0, parsed, available - parsed)
                buffer.seek(0, os.SEEK_SET)
                available -= parsed
                if needed > self.max_packet_size:
                    parser.throw(BufferTooSmallError(
                        needed,
                        self.max_packet_size,
                        offset=offset,
                    ))
                # Leave space for a trailing zero byte
                size, ancdata, msg_flags, _ = request.recvmsg_into(
                    (memoryview(buffer)[available:-1],),
                    socket.CMSG_LEN(3 * SIZEOF_INT),
                    socket.MSG_CMSG_CLOEXEC,
                )
                available += size
                # Ensure that a trailing zero byte exists
                buffer[available] = 0

                streams.extend(
                    stack.enter_context(stream)
                    for stream in self.parse_ancillary_data(ancdata)
                )

                if msg_flags & socket.MSG_CTRUNC:
                    raise ProtocolError("Truncated ancillary data")

                try:
                    needed = parser.send((buffer, available))
                except StopIteration as e:
                    _, _, (argv, environ) = e.value
                    if buffer.tell() < available:
                        raise ProtocolError(
                            "{} byte(s) left over after parsing"
                            .format(available - buffer.tell())
                        )
                    needed = 0
                except BaseParseError as e:
                    raise e.with_offset(offset + buffer.tell())
                else:
                    # Remote end closed/shut down writing
                    if needed > 0 and size == 0:
                        # Raise an error in the parser to produce an error with
                        # proper message
                        parser.throw(UnexpectedEOFError(
                            needed,
                            available,
                            offset=offset + buffer.tell(),
                        ))

            if not streams:
                raise ProtocolError("No file descriptors received")
            if len(streams) != 3:
                raise ProtocolError(
                    "Expected to receive exactly 3 file descriptors"
                )
            stdin = ensure_stream_byte_readable(streams[0], 'stdin')
            stdout = ensure_stream_byte_writable(streams[1], 'stdout')
            stderr = ensure_stream_byte_writable(streams[2], 'stderr')
            # Clear the stack
            stack.pop_all()
            return (stdin, stdout, stderr), argv, environ

    @staticmethod
    def parse_ancillary_data(
            ancdata: Container[Tuple[int, int, bytes]],
    ) -> List[FileIO]:
        """
        Open streams for file descriptors received via :func:`socket.recvmsg`
        ancillary data.

        :param ancdata:
        :return:
        """
        fds = array.array("i")
        streams = []
        with contextlib.ExitStack() as stack:
            truncated = False
            for cmsg_level, cmsg_type, cmsg_data in ancdata:
                if (
                        cmsg_level == socket.SOL_SOCKET
                        and cmsg_type == socket.SCM_RIGHTS
                ):
                    end = len(cmsg_data) - (len(cmsg_data) % fds.itemsize)
                    truncated |= end != len(cmsg_data)
                    fds.frombytes(cmsg_data[:end])
                else:
                    logger.warning(
                        "Received unexpected control message: level=%d type=%d",
                        cmsg_level, cmsg_type,
                    )
            # Ensure that file descriptors get closed on error
            for fd in fds:
                stack.callback(_try_close, fd)
            if truncated:
                raise ProtocolError(
                    "Received truncated file descriptor. "
                    "SCM_RIGHTS control message data must be an multiple of "
                    "sizeof(int) = {}".format(fds.itemsize)
                )
            for fd in fds:
                flags = fcntl.fcntl(fd, fcntl.F_GETFL)
                if flags & os.O_ACCMODE == os.O_RDONLY:
                    mode = "rb"
                elif flags & os.O_ACCMODE == os.O_WRONLY:
                    mode = "wb"
                elif flags & os.O_ACCMODE == os.O_RDWR:
                    # the stdout/stderr buffers can possibly be in RW mode,
                    # however the buffer used by `sys.stdout` is usually opened in
                    # write-only mode by python.
                    # in fact, opening this in `r+` (i.e. read-write mode) and using buffering
                    # causes open() to refuse operation because the buffer is not seekable.
                    # See https://bugs.python.org/issue20074#msg207012 and the related discussion
                    # for some details on the core developers' philosophy on this.
                    mode = "wb"
                else:
                    os.close(fd)
                    continue
                # noinspection PyTypeChecker
                try:
                    stream: FileIO = os.fdopen(fd, mode, buffering=0, closefd=True)
                except io.UnsupportedOperation as e:
                    raise RuntimeError(
                        f"Unable to open fd {fd} ({(len(streams))} already parsed, {mode=})"
                    ) from e
                streams.append(stream)
            stack.pop_all()
            return streams

    @staticmethod
    def parse_int(
            data: mmap.mmap,
            size: int,
            element: str = "int",
    ) -> Parser[int]:
        """Try to parse a C int"""
        need = SIZEOF_INT
        if data.tell() + need > size:
            try:
                data, size = yield need
            except BaseParseError as e:
                raise e.with_element(element)
        value = struct.unpack("=i", data.read(need))[0]
        return data, size, value

    @staticmethod
    def parse_string(
            data: mmap.mmap,
            size: int,
            element: str = "string"
    ) -> Parser[str]:
        """Try to parse a zero-terminated C string"""
        need = 1
        while True:
            if data.tell() + need > size:
                try:
                    data, size = yield need
                except BaseParseError as e:
                    raise e.with_element(element)
            # This is safe, because we ensure that underlying buffer is always
            # zero-terminated
            start = data.tell()
            end = data.find(b'\x00', start, size)
            if end != -1:
                arg = data[start:end]
                data.seek(end + 1, os.SEEK_SET)
                return data, size, arg
            else:
                need = size - start + 1

    @classmethod
    def parse_request(
            cls,
            data: mmap.mmap,
            size: int,
    ) -> Parser[Tuple[List[bytes], Dict[bytes, bytes]]]:
        # Parse number of arguments
        element = "argc"
        data, size, argc = yield from cls.parse_int(data, size, element)
        if argc < 0:
            raise ParseError("Negative argc: " + str(argc), element=element)

        # Parse arguments
        argv = []
        for i in range(argc):
            element = "argv[{:d}]".format(i)
            data, size, arg = yield from cls.parse_string(data, size, element)
            argv.append(arg)

        # Parse number of environment variables
        element = "envc"
        data, size, envc = yield from cls.parse_int(data, size, element)
        if envc < 0:
            raise ParseError("Negative envc: " + str(envc), element=element)

        # Parse environment variables
        environ = {}
        for i in range(envc):
            element = "environ[{}]".format(i)
            data, size, env = yield from cls.parse_string(data, size, element)
            name, sep, value = env.partition(b'=')
            if not sep:
                raise ParseError(
                    "No equal sign in environment variable: " + repr(name),
                    element=element,
                )
            environ[name] = value
        return data, size, (argv, environ)

    def _handle_shutdown_signal(self, signo, _frame):
        logger.error("Received signal %d. Shutting down.", signo)
        # shutdown blocks until the server is stopped, therefore we must use a
        # separate thread, otherwise there will be deadlock
        threading.Thread(name='shutdown', target=self.shutdown).start()

    def serve_forever(self, poll_interval=0.5):
        logger.info("Starting server loop")
        with install_handler(
            (signal.SIGHUP, signal.SIGINT, signal.SIGTERM),
            self._handle_shutdown_signal
        ):
            super().serve_forever(poll_interval)

    @staticmethod
    def _process(
            stdin: FileIO, stdout: FileIO, stderr: FileIO,
            args: Sequence[bytes], env: Dict[bytes, bytes]
    ) -> int:
        # TODO delegate stdout, stderr, etc. to dhcp_script
        return main(
            argv=[a.decode() for a in args],
            standalone=False,
        )

def ensure_stream_byte_readable(stream, stream_desc: str):
    if 'r' not in stream.mode or 'b' not in stream.mode:
        raise ProtocolError(
            f"stream {stream_desc} is not byte-readable ({stream.mode=})"
        )
    return stream

def ensure_stream_byte_writable(stream, stream_desc: str):
    is_writable = bool({'w', '+'} & set(stream.mode))
    if not is_writable or 'b' not in stream.mode:
        raise ProtocolError(
            f"Stream {stream_desc} is not byte-writable ({stream.mode=})"
        )
    return stream

def _try_close(fd):
    try:
        os.close(fd)
    except OSError:
        pass
<commit_msg>Open FDs in text mode instead of binary<commit_after>import array
import contextlib
import ctypes
import fcntl
import io
import logging
import mmap
import os
import signal
import socket
import socketserver
import struct
import sys
import threading
from io import FileIO

from collections.abc import Container
from typing import Dict, Generator, List, Optional, Sequence, Tuple, TypeVar, TextIO

from hades.bin.dhcp_script import main
from hades.common.signals import install_handler

logger = logging.getLogger(__name__)
SIZEOF_INT = ctypes.sizeof(ctypes.c_int)
memfd_create = ctypes.cdll.LoadLibrary("libc.so.6").memfd_create
MFD_CLOEXEC = 1
T = TypeVar('T')
Parser = Generator[int, Tuple[mmap.mmap, int], T]

class BaseParseError(Exception):
    def __init__(
            self,
            *args,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.element = element
        self.offset = offset
        super().__init__(*args, element, offset)

    def _prefix(self) -> str:
        offset, element = self.offset, self.element
        return "".join([
            "" if offset is None else "at offset {}: ".format(offset),
            "" if element is None else "while parsing {}: ".format(element),
        ]).capitalize()

    def with_element(self, element: str):
        self.element = element
        return self

    def with_offset(self, offset: int):
        self.offset = offset
        return self

class ParseError(BaseParseError):
    def __init__(
            self,
            message: str,
            *,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.element = element
        self.offset = offset
        self.message = message
        super().__init__(message, element=element, offset=offset)

    def __str__(self) -> str:
        return self._prefix() + self.message

class UnexpectedEOFError(BaseParseError):
    def __init__(
            self,
            needed: int,
            available: int,
            *,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.needed = needed
        self.available = available
        super().__init__(
            needed,
            available,
            element=element,
            offset=offset,
        )

    def __str__(self) -> str:
        return (
            "{}Unexpected end of file, expected at least {} more byte(s), "
            "but only {} byte(s) left."
            .format(self._prefix(), self.needed, self.available)
        )

class BufferTooSmallError(BaseParseError):
    def __init__(
            self,
            needed: int,
            available: int,
            *,
            element: Optional[str] = None,
            offset: Optional[int] = None,
    ) -> None:
        self.needed = needed
        self.available = available
        super().__init__(
            needed,
            available,
            element=element,
            offset=offset,
        )

    def __str__(self) -> str:
        return (
            "{}Parser requires more data ({}) than can be buffered ({})."
            .format(self._prefix(), self.needed, self.available)
        )

class ProtocolError(Exception):
    pass

class Server(socketserver.UnixStreamServer):
    """
    Process :program:`dnsmasq` :option:`--dhcp-script` invocations.

    :program:`dnsmasq` can notify external tools about DHCP lease activity and
    even maintain a completely external lease database, if additionally
    :option:`--leasefile-ro` is specified. We use this mechanism to store the
    leases in the database.

    Starting a Python script that connects to a PostgreSQL database for DHCP
    lease activity of dnsmasq is too slow however:

    * The Python interpreter has some startup latency, but more importantly
      setuptools console scripts have a very long startup time (>1s) due to
      `this issue <https://github.com/pypa/setuptools/issues/510>`_.
    * PostgreSQL forks for every connection.

    To alleviate these issues we use a small and quick C client that passes its
    command-line arguments, environment variables and file descriptors over a
    UNIX socket to a long running Python server application that's permanently
    connected to the database. The server is single threaded and will only
    handle a single request at a time, because :program:`dnsmasq` itself is
    single threaded.

    The protocol is as follows:

    * We use a SOCK_STREAM AF_UNIX socket, so that the client knows on the
      socket level, when server has received all data and finished handling the
      request.
    * At first the client sends all data to server. The data may not exceed
      :py:code:`mmap.PAGESIZE - 1`. The content of the data is as follows:
        1. The first bytes are the number of arguments (:c:data:`argc`) as
           native type :c:type:`int`.
        2. This is followed by the contents of :c:data:`argv` as a series of
           zero terminated strings.
        3. After that the number of environment variables, that start with the
           prefix :envvar:`DNSMASQ_*` as native type :c:type:`int`.
        4. Followed by the actual environment variables as a series of zero
           terminated strings.
    * The three standard file descriptors stdin, stdout, and stderr of the
      script should be passed with data at some point in time via a
      :c:data:`SCM_RIGHTS` control message.
    * After all data has been sent, the client shuts down its write end of the
      connection and signals the server thereby, that it can begin processing
      the message.
    * The server will process the message and if necessary read additional data
      from the passed stdin file descriptor and write data to the passed stdout
      and stderr file descriptors.
    * After it has processed the message, the server sends a single byte status
      code between 0 and 127 and will close the connection. The script will
      exit with the status code.
    """
    max_packet_size = mmap.PAGESIZE - 1

    def __init__(self, sock, engine):
        self.engine = engine
        server_address = sock.getsockname()
        super().__init__(
            server_address, self._request_handler, bind_and_activate=False,
        )
        self.socket = sock
        # Leave one byte extra for trailing zero byte
        # TODO: With Python 3.8 a memfd can be opened and mapped twice:
        # writable and readonly
        fd = memfd_create(b"buffer", MFD_CLOEXEC)
        os.ftruncate(fd, self.max_packet_size + 1)
        self.buffer = mmap.mmap(
            fd,
            self.max_packet_size + 1,
            mmap.MAP_PRIVATE,
            mmap.PROT_READ | mmap.PROT_WRITE,
        )

    def _request_handler(
            self,
            request: socket.socket,
            client_address,
            server: socketserver.BaseServer,
    ):
        assert self == server

        logger.debug("Received new request from %s", client_address)
        (stdin, stdout, stderr), args, env = self._receive(request)
        status = os.EX_SOFTWARE
        try:
            status = self._process(stdin, stdout, stderr, args, env)
        finally:
            stdout.flush()
            stderr.flush()
            request.send(status.to_bytes(1, sys.byteorder))

    def _receive(
            self,
            request: socket.socket,
    ) -> Tuple[Tuple[TextIO, TextIO, TextIO], List[bytes], Dict[bytes, bytes]]:
        streams: List[TextIO] = []
        # Offset of the buffer relative to the input stream
        offset = 0
        # Number of filled bytes in the buffer
        available = 0
        # Initialize parser
        buffer = self.buffer
        buffer.seek(0, os.SEEK_SET)
        parser = self.parse_request(buffer, 0)
        needed = next(parser)
        with contextlib.ExitStack() as stack:
            while needed:
                # Prepare buffer for refill
                parsed = buffer.tell()
                offset += parsed
                buffer.move(0, parsed, available - parsed)
                buffer.seek(0, os.SEEK_SET)
                available -= parsed
                if needed > self.max_packet_size:
                    parser.throw(BufferTooSmallError(
                        needed,
                        self.max_packet_size,
                        offset=offset,
                    ))
                # Leave space for a trailing zero byte
                size, ancdata, msg_flags, _ = request.recvmsg_into(
                    (memoryview(buffer)[available:-1],),
                    socket.CMSG_LEN(3 * SIZEOF_INT),
                    socket.MSG_CMSG_CLOEXEC,
                )
                available += size
                # Ensure that a trailing zero byte exists
                buffer[available] = 0

                streams.extend(
                    stack.enter_context(stream)
                    for stream in self.parse_ancillary_data(ancdata)
                )

                if msg_flags & socket.MSG_CTRUNC:
                    raise ProtocolError("Truncated ancillary data")

                try:
                    needed = parser.send((buffer, available))
                except StopIteration as e:
                    _, _, (argv, environ) = e.value
                    if buffer.tell() < available:
                        raise ProtocolError(
                            "{} byte(s) left over after parsing"
                            .format(available - buffer.tell())
                        )
                    needed = 0
                except BaseParseError as e:
                    raise e.with_offset(offset + buffer.tell())
                else:
                    # Remote end closed/shut down writing
                    if needed > 0 and size == 0:
                        # Raise an error in the parser to produce an error with
                        # proper message
                        parser.throw(UnexpectedEOFError(
                            needed,
                            available,
                            offset=offset + buffer.tell(),
                        ))

            if not streams:
                raise ProtocolError("No file descriptors received")
            if len(streams) != 3:
                raise ProtocolError(
                    "Expected to receive exactly 3 file descriptors"
                )
            stdin = ensure_stream_readable(streams[0], 'stdin')
            stdout = ensure_stream_writable(streams[1], 'stdout')
            stderr = ensure_stream_writable(streams[2], 'stderr')
            # Clear the stack
            stack.pop_all()
            return (stdin, stdout, stderr), argv, environ

    @staticmethod
    def parse_ancillary_data(
            ancdata: Container[Tuple[int, int, bytes]],
    ) -> List[FileIO]:
        """
        Open streams for file descriptors received via :func:`socket.recvmsg`
        ancillary data.

        :param ancdata:
        :return:
        """
        fds = array.array("i")
        streams = []
        with contextlib.ExitStack() as stack:
            truncated = False
            for cmsg_level, cmsg_type, cmsg_data in ancdata:
                if (
                        cmsg_level == socket.SOL_SOCKET
                        and cmsg_type == socket.SCM_RIGHTS
                ):
                    end = len(cmsg_data) - (len(cmsg_data) % fds.itemsize)
                    truncated |= end != len(cmsg_data)
                    fds.frombytes(cmsg_data[:end])
                else:
                    logger.warning(
                        "Received unexpected control message: level=%d type=%d",
                        cmsg_level, cmsg_type,
                    )
            # Ensure that file descriptors get closed on error
            for fd in fds:
                stack.callback(_try_close, fd)
            if truncated:
                raise ProtocolError(
                    "Received truncated file descriptor. "
                    "SCM_RIGHTS control message data must be an multiple of "
                    "sizeof(int) = {}".format(fds.itemsize)
                )
            for fd in fds:
                flags = fcntl.fcntl(fd, fcntl.F_GETFL)
                if flags & os.O_ACCMODE == os.O_RDONLY:
                    mode = "r"
                elif flags & os.O_ACCMODE == os.O_WRONLY:
                    mode = "w"
                elif flags & os.O_ACCMODE == os.O_RDWR:
                    # the stdout/stderr buffers can possibly be in RW mode,
                    # however the buffer used by `sys.stdout` is usually opened in
                    # write-only mode by python.
                    # in fact, opening this in `r+` (i.e. read-write mode) and using buffering
                    # causes open() to refuse operation because the buffer is not seekable.
                    # See https://bugs.python.org/issue20074#msg207012 and the related discussion
                    # for some details on the core developers' philosophy on this.
                    mode = "wb"
                else:
                    os.close(fd)
                    continue
                # noinspection PyTypeChecker
                try:
                    stream: TextIO = os.fdopen(fd, mode, closefd=True)
                except io.UnsupportedOperation as e:
                    raise RuntimeError(
                        f"Unable to open fd {fd} ({(len(streams))} already parsed, {mode=})"
                    ) from e
                streams.append(stream)
            stack.pop_all()
            return streams

    @staticmethod
    def parse_int(
            data: mmap.mmap,
            size: int,
            element: str = "int",
    ) -> Parser[int]:
        """Try to parse a C int"""
        need = SIZEOF_INT
        if data.tell() + need > size:
            try:
                data, size = yield need
            except BaseParseError as e:
                raise e.with_element(element)
        value = struct.unpack("=i", data.read(need))[0]
        return data, size, value

    @staticmethod
    def parse_string(
            data: mmap.mmap,
            size: int,
            element: str = "string"
    ) -> Parser[str]:
        """Try to parse a zero-terminated C string"""
        need = 1
        while True:
            if data.tell() + need > size:
                try:
                    data, size = yield need
                except BaseParseError as e:
                    raise e.with_element(element)
            # This is safe, because we ensure that underlying buffer is always
            # zero-terminated
            start = data.tell()
            end = data.find(b'\x00', start, size)
            if end != -1:
                arg = data[start:end]
                data.seek(end + 1, os.SEEK_SET)
                return data, size, arg
            else:
                need = size - start + 1

    @classmethod
    def parse_request(
            cls,
            data: mmap.mmap,
            size: int,
    ) -> Parser[Tuple[List[bytes], Dict[bytes, bytes]]]:
        # Parse number of arguments
        element = "argc"
        data, size, argc = yield from cls.parse_int(data, size, element)
        if argc < 0:
            raise ParseError("Negative argc: " + str(argc), element=element)

        # Parse arguments
        argv = []
        for i in range(argc):
            element = "argv[{:d}]".format(i)
            data, size, arg = yield from cls.parse_string(data, size, element)
            argv.append(arg)

        # Parse number of environment variables
        element = "envc"
        data, size, envc = yield from cls.parse_int(data, size, element)
        if envc < 0:
            raise ParseError("Negative envc: " + str(envc), element=element)

        # Parse environment variables
        environ = {}
        for i in range(envc):
            element = "environ[{}]".format(i)
            data, size, env = yield from cls.parse_string(data, size, element)
            name, sep, value = env.partition(b'=')
            if not sep:
                raise ParseError(
                    "No equal sign in environment variable: " + repr(name),
                    element=element,
                )
            environ[name] = value
        return data, size, (argv, environ)

    def _handle_shutdown_signal(self, signo, _frame):
        logger.error("Received signal %d. Shutting down.", signo)
        # shutdown blocks until the server is stopped, therefore we must use a
        # separate thread, otherwise there will be deadlock
        threading.Thread(name='shutdown', target=self.shutdown).start()

    def serve_forever(self, poll_interval=0.5):
        logger.info("Starting server loop")
        with install_handler(
            (signal.SIGHUP, signal.SIGINT, signal.SIGTERM),
            self._handle_shutdown_signal
        ):
            super().serve_forever(poll_interval)

    @staticmethod
    def _process(
            stdin: TextIO, stdout: TextIO, stderr: TextIO,
            args: Sequence[bytes], env: Dict[bytes, bytes]
    ) -> int:
        # TODO delegate stdout, stderr, etc. to dhcp_script
        return main(
            argv=[a.decode() for a in args],
            standalone=False,
        )

def ensure_stream_readable(stream, stream_desc: str):
    if 'r' not in stream.mode:
        raise ProtocolError(
            f"stream {stream_desc} is not readable ({stream.mode=})"
        )
    return stream

def ensure_stream_writable(stream, stream_desc: str):
    is_writable = bool({'w', '+'} & set(stream.mode))
    if not is_writable:
        raise ProtocolError(
            f"Stream {stream_desc} is not writable ({stream.mode=})"
        )
    return stream

def _try_close(fd):
    try:
        os.close(fd)
    except OSError:
        pass
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
Perform data manipulation tasks and create inputs for project workflow
"""

import os
try:
    import urllib.request as urllib2
except ImportError:
    import urllib2

import boto
import numpy as np
import pandas as pd
import sklearn
import boto.s3
from boto.s3.key import Key
from pychem import getmol
from sklearn.ensemble import RandomForestRegressor

__author__ = "Pearl Philip"
__credits__ = "David Beck"
__license__ = "BSD 3-Clause License"
__maintainer__ = "Pearl Philip"
__email__ = "pphilip@uw.edu"
__status__ = "Development"

def create_notation_dataframe(filename):
    """
    Returns Pandas dataframe of sample ID and molecular notation

    Inputs: filename
    filename - path to file containing molecular notation indexed by sample ID
    Input types: str

    Outputs: dataframe of molecular notation indexed by sample ID
    Output types: Pandas DataFrame
    """
    df = []
    for line in filename:
        # Splits the line into it's key and molecular string
        words = line.split()
        z = [int(words[0]), words[1]]
        df.append(z)
    df = pd.DataFrame(df)
    df.columns = ['CID', 'SMILES']
    df.sort_values(by='CID', inplace=True)
    return df

def create_activity_dataframe(dataframe):
    """
    Performing useful transformations on the acquired data for use in subsequent algorithm.
    :param dataframe: Dataframe downloaded from NCBI database.
    :return: df: Cleaned and sorted dataframe.
    """

    # Eliminates first five text rows of csv
    for j in range(5):
        df = dataframe.drop(j, axis=0)
    df = df.drop(['PUBCHEM_ACTIVITY_URL', 'PUBCHEM_RESULT_TAG',
                  'PUBCHEM_ACTIVITY_SCORE', 'PUBCHEM_SID',
                  'PUBCHEM_ASSAYDATA_COMMENT', 'Potency',
                  'Efficacy', 'Analysis Comment',
                  'Curve_Description', 'Fit_LogAC50',
                  'Fit_HillSlope', 'Fit_R2', 'Fit_ZeroActivity',
                  'Fit_CurveClass', 'Excluded_Points', 'Compound QC',
                  'Max_Response', 'Phenotype', 'Activity at 0.457 uM',
                  'Activity at 2.290 uM', 'Activity at 11.40 uM',
                  'Activity at 57.10 uM', 'PUBCHEM_ACTIVITY_OUTCOME',
                  'Fit_InfiniteActivity'], axis=1)
    df.rename(columns={'PUBCHEM_CID': 'CID'}, inplace=True)

    # Eliminates duplicate compound rows
    df['dupes'] = df.duplicated('CID')
    df = df[df['dupes'] == 0].drop(['dupes'], axis=1)
    df = df.sort_values(by='CID')
    df.drop(276743)
    """"# Extract SMILES strings from CID values
    smiles = []
    i = 0
    for cid in df['CID']:
        string = getmol.GetMolFromNCBI(cid='%d' % cid)
        smiles.append(string)
        print(i)
        i += 1
    df['SMILES'] = smiles
    df.to_csv('data/NCBIdata.csv')"""
    return df

def upload_to_s3(aws_access_key_id, aws_secret_access_key, file_to_s3, bucket, key, callback=None, md5=None,
                 reduced_redundancy=False, content_type=None):
    """
    Uploads the given file to the AWS S3 bucket and key specified.

    callback is a function of the form:

    def callback(complete, total)

    The callback should accept two integer parameters, the first representing the number of bytes that have been
    successfully transmitted to S3 and the second representing the size of the to be transmitted object.

    Returns boolean indicating success/failure of upload.
    """
    try:
        size = os.fstat(file_to_s3.fileno()).st_size
    except:
        # Not all file objects implement fileno(),
        # so we fall back on this
        file_to_s3.seek(0, os.SEEK_END)
        size = file_to_s3.tell()

    conn = boto.connect_s3(aws_access_key_id, aws_secret_access_key)
    bucket = conn.get_bucket(bucket, validate=True)
    k = Key(bucket)
    k.key = key
    if content_type:
        k.set_metadata('Content-Type', content_type)
    sent = k.set_contents_from_file(file_to_s3, cb=callback, md5=md5, reduced_redundancy=reduced_redundancy,
                                    rewind=True)

    # Rewind for later use
    file_to_s3.seek(0)

    if sent == size:
        return True
    return False

def join_dataframes():
    """
    Joins the dataframes of existing descriptor files from their urls into a single dataframe.
    :return: Dataframe after joining
    """

    url_list = ['https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_constitution.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_con.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_kappa.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_estate.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_basak.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_property.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_charge.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_moe.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_burden.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_geary.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_moran.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_topology.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_geometric.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_cpsa.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_rdf.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_morse.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_whim.csv'
                ]

    url_exist_list = []
    for url in url_list:
        try:
            r = urllib2.urlopen(url)
        except urllib2.URLError as e:
            r = e
        if r.code < 400:
            url_exist_list.append(url)
        else:
            None

    i = 0
    df = [0] * len(url_exist_list)
    for url in url_exist_list:
        df[i] = pd.read_csv(url)

        df[i].drop(df[i].columns[0], axis=1, inplace=True)
        df[i].reset_index(drop=True, inplace=True)
        i += 1

    joined_df = df[0]
    for i in df[1:]:
        joined_df = joined_df.join(i)
    return joined_df

def choose_features(x, y):
    """
    Selecting the features of high importance to reduce feature space.
    :param x: Dataframe of features
    :param y: Dataframe of target property
    :return desired x: Dataframe of short-listed features
    """

    # Random forest feature importance
    x = np.array(x)
    y = np.array(y)

    clf = RandomForestRegressor()
    sfm = sklearn.feature_selection.SelectFromModel(clf, threshold=0.15)
    sfm.fit(x, y)
    print 'here'
    desired_x = sfm.transform(x)

    return desired_x

def remove_nan_infinite(dataframe):
    """
    :param dataframe: Dataframe containing NaN and infinite values
    :return dataframe: Dataframe with no NaN or infinite values
    """

    dataframe.replace([np.inf, -np.inf], np.nan, inplace=True)
    dataframe.fillna(0, inplace=True)

    return dataframe

<commit_msg>Updated all docstrings in utils.py<commit_after>#!/usr/bin/env python

"""
Create inputs and perform data manipulation tasks in project workflow
"""

import os
try:
    import urllib.request as urllib2
except ImportError:
    import urllib2

import boto
import numpy as np
import pandas as pd
import sklearn
import boto.s3
from boto.s3.key import Key
from pychem import getmol
from sklearn.ensemble import RandomForestRegressor

__author__ = "Pearl Philip"
__credits__ = "David Beck"
__license__ = "BSD 3-Clause License"
__maintainer__ = "Pearl Philip"
__email__ = "pphilip@uw.edu"
__status__ = "Development"

def create_notation_dataframe(filename):
    """
    Returning Pandas dataframe of sample ID and molecular notation.
    :param filename: File object containing molecular notation indexed by sample ID
    :return: Dataframe of molecular notation indexed by sample ID.
    """
    df = []
    for line in filename:
        # Splits the line into it's key and molecular string
        words = line.split()
        z = [int(words[0]), words[1]]
        df.append(z)
    df = pd.DataFrame(df)
    df.columns = ['CID', 'SMILES']
    df.sort_values(by='CID', inplace=True)
    return df

def create_activity_dataframe(dataframe):
    """
    Performing useful transformations on the acquired data for use in subsequent algorithm.
    :param dataframe: Dataframe downloaded from NCBI database.
    :return: df: Cleaned and sorted dataframe.
    """

    # Eliminates first five text rows of csv
    for j in range(5):
        df = dataframe.drop(j, axis=0)
    df = df.drop(['PUBCHEM_ACTIVITY_URL', 'PUBCHEM_RESULT_TAG',
                  'PUBCHEM_ACTIVITY_SCORE', 'PUBCHEM_SID',
                  'PUBCHEM_ASSAYDATA_COMMENT', 'Potency',
                  'Efficacy', 'Analysis Comment',
                  'Curve_Description', 'Fit_LogAC50',
                  'Fit_HillSlope', 'Fit_R2', 'Fit_ZeroActivity',
                  'Fit_CurveClass', 'Excluded_Points', 'Compound QC',
                  'Max_Response', 'Phenotype', 'Activity at 0.457 uM',
                  'Activity at 2.290 uM', 'Activity at 11.40 uM',
                  'Activity at 57.10 uM', 'PUBCHEM_ACTIVITY_OUTCOME',
                  'Fit_InfiniteActivity'], axis=1)
    df.rename(columns={'PUBCHEM_CID': 'CID'}, inplace=True)

    # Eliminates duplicate compound rows
    df['dupes'] = df.duplicated('CID')
    df = df[df['dupes'] == 0].drop(['dupes'], axis=1)
    df = df.sort_values(by='CID')
    df.drop(276743)
    """"# Extract SMILES strings from CID values
    smiles = []
    i = 0
    for cid in df['CID']:
        string = getmol.GetMolFromNCBI(cid='%d' % cid)
        smiles.append(string)
        print(i)
        i += 1
    df['SMILES'] = smiles
    df.to_csv('data/NCBIdata.csv')"""
    return df

def upload_to_s3(aws_access_key_id, aws_secret_access_key, file_to_s3, bucket, key, callback=None, md5=None,
                 reduced_redundancy=False, content_type=None):
    """
    Uploads the given file to the AWS S3 bucket and key specified.
    :param aws_access_key_id: First part of AWS access key.
    :param aws_secret_access_key: Second part of AWS access key.
    :param file_to_s3: File object to be uploaded.
    :param bucket: S3 bucket name as string.
    :param key: Name attribute of the file object to be uploaded.
    :param callback: Function accepts two integer parameters, the first representing the number of bytes that have been
    successfully transmitted to S3 and the second representing the size of the to be transmitted object. Returns
    boolean indicating success/failure of upload.
    :param md5: MD5 checksum value to verify the integrity of the object.
    :param reduced_redundancy: S3 option that enables customers to reduce their costs
    by storing noncritical, reproducible data at lower levels of redundancy than S3's standard storage.
    :param content_type: Set the type of content in file object.
    :return: Boolean indicating success of upload.
    """
    try:
        size = os.fstat(file_to_s3.fileno()).st_size
    except:
        # Not all file objects implement fileno(), so we fall back on this
        file_to_s3.seek(0, os.SEEK_END)
        size = file_to_s3.tell()

    conn = boto.connect_s3(aws_access_key_id, aws_secret_access_key)
    bucket = conn.get_bucket(bucket, validate=True)
    k = Key(bucket)
    k.key = key
    if content_type:
        k.set_metadata('Content-Type', content_type)
    sent = k.set_contents_from_file(file_to_s3, cb=callback, md5=md5,
                                    reduced_redundancy=reduced_redundancy, rewind=True)

    # Rewind for later use
    file_to_s3.seek(0)

    if sent == size:
        return True
    return False

def join_dataframes():
    """
    Joining the dataframes of existing descriptor files from their urls into a single dataframe.
    :return: Dataframe after join over key column.
    """

    url_list = ['https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_constitution.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_con.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_kappa.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_estate.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_basak.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_property.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_charge.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_moe.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_burden.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_geary.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_moran.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_topology.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_geometric.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_cpsa.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_rdf.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_morse.csv',
                'https://s3-us-west-2.amazonaws.com/pphilip-usp-inhibition/df_whim.csv'
                ]

    url_exist_list = []
    for url in url_list:
        try:
            r = urllib2.urlopen(url)
        except urllib2.URLError as e:
            r = e
        if r.code < 400:
            url_exist_list.append(url)

    i = 0
    df = [0] * len(url_exist_list)
    for url in url_exist_list:
        df[i] = pd.read_csv(url)

        df[i].drop(df[i].columns[0], axis=1, inplace=True)
        df[i].reset_index(drop=True, inplace=True)
        i += 1

    joined_df = df[0]
    for i in df[1:]:
        joined_df = joined_df.join(i)
    return joined_df

def choose_features(x, y):
    """
    Selecting the features of high importance to reduce feature space.
    :param x: Dataframe of features.
    :param y: Dataframe of target property.
    :return desired x: Dataframe of short-listed features.
    """

    # Random forest feature importance
    x = np.array(x)
    y = np.array(y)

    clf = RandomForestRegressor()
    sfm = sklearn.feature_selection.SelectFromModel(clf, threshold=0.15)
    sfm.fit(x, y)
    desired_x = sfm.transform(x)

    return desired_x

def remove_nan_infinite(dataframe):
    """
    Removing NaN and infinite values from the dataframe.
    :param dataframe: Dataframe containing NaN and infinite values.
    :return dataframe: Dataframe with no NaN or infinite values.
    """

    dataframe.replace([np.inf, -np.inf], np.nan, inplace=True)
    dataframe.fillna(0, inplace=True)

    return dataframe
<|endoftext|>


<commit_before>import numpy as np

class SpheroidLocation(object):
    """ Class representing a ground station on an arbitrary ellipsoid.
    """

    def __init__(self, lon, lat, h, body):
        """

        Parameters
        ----------
        lon: ~astropy.units.quantity.Quantity
            geodetic longitude

        lat: ~astropy.units.quantity.Quantity
            geodetic latitude

        h: ~astropy.units.quantity.Quantity
            geodetic height

        body: ~poliastro.bodies.Body
            Planetary body the spheroid location lies on

        """
        self._lon = lon
        self._lat = lat
        self._h = h
        self._body = body
        self._a = body.R
        self._b = body.R
        self._c = body.R_polar

    @property
    def cartesian_cords(self):
        """Convert to the Cartesian Coordinate system.
        """
        e2 = 1 - (self._c / self._a) ** 2
        N = self._a / np.sqrt(1 - e2 * np.sin(self._lon) ** 2)

        x = (N + self._h) * np.cos(self._lon) * np.cos(self._lat)
        y = (N + self._h) * np.cos(self._lon) * np.sin(self._lat)
        z = ((1 - e2) * N + self._h) * np.sin(self._lon)
        return x, y, z

    @property
    def f(self):
        """Get first flattening.
        """
        return 1 - self._c / self._a

    @property
    def N(self):
        """Normal vector of the ellipsoid at the given location.
        """
        x, y, z = self.cartesian_cords
        a, b, c = self._a.value, self._b.value, self._c.value
        N = np.array([2 * x.value / a ** 2, 2 * y.value / b ** 2, 2 * z.value / c ** 2])
        N /= np.linalg.norm(N)
        return N

    @property
    def tangential_vecs(self):
        """Returns orthonormal vectors tangential to the ellipsoid at the given location.
        """
        N = self.N
        u = np.array([1.0, 0, 0])
        u -= u.dot(N) * N
        u /= np.linalg.norm(u)
        v = np.cross(N, u)
        return u, v

    @property
    def radius_of_curvature(self):
        """Radius of curvature of the meridian at the latitude of the given location.
        """
        e2 = 1 - (self._c / self._a) ** 2
        rc = self._a * (1 - e2) / (1 - e2 * np.sin(self._lat) ** 2) ** 1.5
        return rc

    def distance(self, px, py, pz):
        """
        Calculates the distane from an arbitrary point to the given location (Cartesian coordinates).

        Parameters
        ----------
        px: ~astropy.units.quantity.Quantity
            x-coordinate of the point
        py: ~astropy.units.quantity.Quantity
            y-coordinate of the point
        pz: ~astropy.units.quantity.Quantity
            z-coordinate of the point
        """
        c = np.array([c.value for c in self.cartesian_cords])
        u = np.array([px.value, py.value, pz.value])
        d = np.linalg.norm(c - u)
        return d

    def is_visible(self, px, py, pz):
        """
        Determines whether an object located at a given point is visible from the given location.
        Returns true if true, false otherwise.

        Parameters
        ----------
        px: ~astropy.units.quantity.Quantity
            x-coordinate of the point
        py: ~astropy.units.quantity.Quantity
            y-coordinate of the point
        pz: ~astropy.units.quantity.Quantity
            z-coordinate of the point
        """
        c = [c.value for c in self.cartesian_cords]
        u = np.array([px.value, py.value, pz.value])

        N = self.N
        d = -N.dot(c)
        p = N.dot(u) + d

        return p >= 0

    def cartesian_to_ellipsoidal(self, x, y, z):
        """
        Converts ellipsoidal coordinates to the Cartesian coordinate system for the given ellipsoid.
        Instead of the iterative formula, the function uses the approximation introduced in
        Bowring, B. R. (1976). TRANSFORMATION FROM SPATIAL TO GEOGRAPHICAL COORDINATES

        Parameters
        ----------
        x: ~astropy.units.quantity.Quantity
            x coordinate

        y: ~astropy.units.quantity.Quantity
            y coordinate

        z: ~astropy.units.quantity.Quantity
            z coordinate

        """
        a = self._a
        c = self._c
        e2 = 1 - (c / a) ** 2
        e2_ = e2 / (1 - e2)
        p = np.sqrt(x ** 2 + y ** 2)
        th = np.arctan(z * a / (p * c))
        lon = np.arctan(y / x)
        lat = np.arctan(
            (z + e2_ * c * np.sin(th) ** 3) / (p - e2 * a * np.cos(th) ** 3)
        )

        v = a / np.sqrt(1 - e2 * np.sin(lat) ** 2)
        h = p / np.cos(lat) - v
        h = z / np.sin(lat) - (1 - e2) * v

        return lat, lon, h
<commit_msg>Remove body as an instance variable<commit_after>import numpy as np

class SpheroidLocation(object):
    """ Class representing a ground station on an arbitrary ellipsoid.
    """

    def __init__(self, lon, lat, h, body):
        """

        Parameters
        ----------
        lon: ~astropy.units.quantity.Quantity
            geodetic longitude

        lat: ~astropy.units.quantity.Quantity
            geodetic latitude

        h: ~astropy.units.quantity.Quantity
            geodetic height

        body: ~poliastro.bodies.Body
            Planetary body the spheroid location lies on

        """
        self._lon = lon
        self._lat = lat
        self._h = h
        self._a = body.R
        self._b = body.R
        self._c = body.R_polar

    @property
    def cartesian_cords(self):
        """Convert to the Cartesian Coordinate system.
        """
        e2 = 1 - (self._c / self._a) ** 2
        N = self._a / np.sqrt(1 - e2 * np.sin(self._lon) ** 2)

        x = (N + self._h) * np.cos(self._lon) * np.cos(self._lat)
        y = (N + self._h) * np.cos(self._lon) * np.sin(self._lat)
        z = ((1 - e2) * N + self._h) * np.sin(self._lon)
        return x, y, z

    @property
    def f(self):
        """Get first flattening.
        """
        return 1 - self._c / self._a

    @property
    def N(self):
        """Normal vector of the ellipsoid at the given location.
        """
        x, y, z = self.cartesian_cords
        a, b, c = self._a.value, self._b.value, self._c.value
        N = np.array([2 * x.value / a ** 2, 2 * y.value / b ** 2, 2 * z.value / c ** 2])
        N /= np.linalg.norm(N)
        return N

    @property
    def tangential_vecs(self):
        """Returns orthonormal vectors tangential to the ellipsoid at the given location.
        """
        N = self.N
        u = np.array([1.0, 0, 0])
        u -= u.dot(N) * N
        u /= np.linalg.norm(u)
        v = np.cross(N, u)
        return u, v

    @property
    def radius_of_curvature(self):
        """Radius of curvature of the meridian at the latitude of the given location.
        """
        e2 = 1 - (self._c / self._a) ** 2
        rc = self._a * (1 - e2) / (1 - e2 * np.sin(self._lat) ** 2) ** 1.5
        return rc

    def distance(self, px, py, pz):
        """
        Calculates the distane from an arbitrary point to the given location (Cartesian coordinates).

        Parameters
        ----------
        px: ~astropy.units.quantity.Quantity
            x-coordinate of the point
        py: ~astropy.units.quantity.Quantity
            y-coordinate of the point
        pz: ~astropy.units.quantity.Quantity
            z-coordinate of the point
        """
        c = np.array([c.value for c in self.cartesian_cords])
        u = np.array([px.value, py.value, pz.value])
        d = np.linalg.norm(c - u)
        return d

    def is_visible(self, px, py, pz):
        """
        Determines whether an object located at a given point is visible from the given location.
        Returns true if true, false otherwise.

        Parameters
        ----------
        px: ~astropy.units.quantity.Quantity
            x-coordinate of the point
        py: ~astropy.units.quantity.Quantity
            y-coordinate of the point
        pz: ~astropy.units.quantity.Quantity
            z-coordinate of the point
        """
        c = [c.value for c in self.cartesian_cords]
        u = np.array([px.value, py.value, pz.value])

        N = self.N
        d = -N.dot(c)
        p = N.dot(u) + d

        return p >= 0

    def cartesian_to_ellipsoidal(self, x, y, z):
        """
        Converts ellipsoidal coordinates to the Cartesian coordinate system for the given ellipsoid.
        Instead of the iterative formula, the function uses the approximation introduced in
        Bowring, B. R. (1976). TRANSFORMATION FROM SPATIAL TO GEOGRAPHICAL COORDINATES

        Parameters
        ----------
        x: ~astropy.units.quantity.Quantity
            x coordinate

        y: ~astropy.units.quantity.Quantity
            y coordinate

        z: ~astropy.units.quantity.Quantity
            z coordinate

        """
        a = self._a
        c = self._c
        e2 = 1 - (c / a) ** 2
        e2_ = e2 / (1 - e2)
        p = np.sqrt(x ** 2 + y ** 2)
        th = np.arctan(z * a / (p * c))
        lon = np.arctan(y / x)
        lat = np.arctan(
            (z + e2_ * c * np.sin(th) ** 3) / (p - e2 * a * np.cos(th) ** 3)
        )

        v = a / np.sqrt(1 - e2 * np.sin(lat) ** 2)
        h = p / np.cos(lat) - v
        h = z / np.sin(lat) - (1 - e2) * v

        return lat, lon, h
<|endoftext|>


<commit_before>#!/usr/bin/env python3
from DGAMalware import DGAMalware

"""TLDs used by Necurs.

"""
tlds = ['tj','in','jp','tw','ac','cm','la','mn','so','sh','sc','nu','nf','mu',
        'ms','mx','ki','im','cx','cc','tv','bz','me','eu','de','ru','co','su','pw',
        'kz','sx','us','ug','ir','to','ga','com','net','org','biz','xxx','pro','bit']

"""Domain Generation Algorithm for Necurs.

Source:
  http://www.johannesbader.ch/2015/02/the-dgas-of-necurs/
  Johannes Bader
  February 20, 2015

"""
class Necurs(DGAMalware):

  @classmethod
  def domainsLifetime(self):
    return 4 * 24 * 3600

  @classmethod
  def domainsFor(self, date):
    domains = []
    for sequence_nr in range(2048):
      domains.append(self.generateDomain(sequence_nr, 9, date))
    return domains

  @classmethod
  def generateDomain(self, sequence_nr, magic_nr, date):
    n = self.pseudoRandom(date.year)
    n = self.pseudoRandom(n + date.month + 43690)
    n = self.pseudoRandom(n + (date.day>>2))
    n = self.pseudoRandom(n + sequence_nr)
    n = self.pseudoRandom(n + magic_nr)
    domain_length = self.mod64(n, 15) + 7

    domain = ""
    for i in range(domain_length):
      n = self.pseudoRandom(n+i)
      ch = self.mod64(n, 25) + ord('a')
      domain += chr(ch)
      n += 0xABBEDF
      n = self.pseudoRandom(n)

    tld = tlds[self.mod64(n, 43)]
    domain += '.' + tld
    return domain

  @classmethod
  def pseudoRandom(self, value):
    loops = (value & 0x7F) + 21
    for index in range(loops):
      value += ((value * 7) ^ (value << 15)) + 8 * index - (value >> 5)
      value &= ((1 << 64) - 1)
    return value

  @classmethod
  def mod64(self, nr1, nr2):
    return nr1 % nr2

  @classmethod
  def couldUseDomain(self, domain):
    useTLD = False
    for tld in tlds:
      if domain.endswith(tld):
        useTLD = True
        break
    if not useTLD:
      return False
    return not any(char.isdigit() for char in domain)
<commit_msg>Improve couldUseDomain method for Necurs<commit_after>#!/usr/bin/env python3
from DGAMalware import DGAMalware

"""TLDs used by Necurs.

"""
tlds = ['tj','in','jp','tw','ac','cm','la','mn','so','sh','sc','nu','nf','mu',
        'ms','mx','ki','im','cx','cc','tv','bz','me','eu','de','ru','co','su','pw',
        'kz','sx','us','ug','ir','to','ga','com','net','org','biz','xxx','pro','bit']

"""Domain Generation Algorithm for Necurs.

Source:
  http://www.johannesbader.ch/2015/02/the-dgas-of-necurs/
  Johannes Bader
  February 20, 2015

"""
class Necurs(DGAMalware):

  @classmethod
  def domainsLifetime(self):
    return 4 * 24 * 3600

  @classmethod
  def domainsFor(self, date):
    domains = []
    for sequence_nr in range(2048):
      domains.append(self.generateDomain(sequence_nr, 9, date))
    return domains

  @classmethod
  def generateDomain(self, sequence_nr, magic_nr, date):
    n = self.pseudoRandom(date.year)
    n = self.pseudoRandom(n + date.month + 43690)
    n = self.pseudoRandom(n + (date.day>>2))
    n = self.pseudoRandom(n + sequence_nr)
    n = self.pseudoRandom(n + magic_nr)
    domain_length = self.mod64(n, 15) + 7

    domain = ""
    for i in range(domain_length):
      n = self.pseudoRandom(n+i)
      ch = self.mod64(n, 25) + ord('a')
      domain += chr(ch)
      n += 0xABBEDF
      n = self.pseudoRandom(n)

    tld = tlds[self.mod64(n, 43)]
    domain += '.' + tld
    return domain

  @classmethod
  def pseudoRandom(self, value):
    loops = (value & 0x7F) + 21
    for index in range(loops):
      value += ((value * 7) ^ (value << 15)) + 8 * index - (value >> 5)
      value &= ((1 << 64) - 1)
    return value

  @classmethod
  def mod64(self, nr1, nr2):
    return nr1 % nr2

  @classmethod
  def couldUseDomain(self, domain):
    useTLD = False
    for tld in tlds:
      if domain.endswith(tld):
        useTLD = True
        break
    if not useTLD:
      return False

    label = domain.split('.', 1)[0]
    if 'z' in label:
      return False

    return not any(char.isdigit() for char in domain)
<|endoftext|>


<commit_before>"""
@file
@brief Helpers to convert docstring to various format

.. versionadded:: 1.0
"""
from ..sphinxext.sphinx_bigger_extension import setup as setup_bigger
from ..sphinxext.sphinx_githublink_extension import setup as setup_githublink
from ..sphinxext.sphinx_blocref_extension import setup as setup_blocref
from ..sphinxext.sphinx_blog_extension import setup as setup_blog
from ..sphinxext.sphinx_exref_extension import setup as setup_exref
from ..sphinxext.sphinx_faqref_extension import setup as setup_faqref
from ..sphinxext.sphinx_mathdef_extension import setup as setup_mathdef
from ..sphinxext.sphinx_nbref_extension import setup as setup_nbref
from ..sphinxext.sphinx_runpython_extension import setup as setup_runpython
from ..sphinxext.sphinx_sharenet_extension import setup as setup_sharenet
from ..sphinxext.sphinx_todoext_extension import setup as setup_todoext
from .convert_doc_sphinx_helper import HTMLWriterWithCustomDirectives
from docutils import nodes
from docutils.parsers.rst import directives as doc_directives, roles as doc_roles
from sphinx.config import Config
from sphinx import __display_version__ as sphinx__display_version__

try:
    from sphinx.util.docutils import is_html5_writer_available
except ImportError:
    # Available only after Sphinx >= 1.6.1.
    def is_html5_writer_available():
        return False

class MockSphinxApp:
    """
    Mock Sphinx application
    """

    def __init__(self, writer, app):
        self.app = app
        self.new_options = {}
        self.writer = writer
        self.mapping = {"<class 'sphinx.ext.todo.todo_node'>": "todo",
                        "<class 'sphinx.ext.graphviz.graphviz'>": "graphviz",
                        "<class 'sphinx.ext.mathbase.math'>": "math",
                        "<class 'sphinx.ext.mathbase.displaymath'>": "displaymath",
                        "<class 'sphinx.ext.mathbase.eqref'>": "eqref",
                        "<class 'matplotlib.sphinxext.only_directives.html_only'>": "only",
                        "<class 'matplotlib.sphinxext.only_directives.latex_only'>": "only",
                        }
        self.mapping_connect = {}
        self.config = Config(None, None, overrides={}, tags=None)
        self.confdir = "."
        self.doctreedir = "."
        self.srcdir = "."
        self.builder = writer.builder
        self.domains = {}
        self._events = {}

    def add_directive(self, name, cl, *args, **options):
        doc_directives.register_directive(name, cl)
        self.mapping[str(cl)] = name
        self.app.add_directive(name, cl, *args, **options)
        self.writer.app.add_directive(name, cl, *args, **options)

    def add_role(self, name, cl):
        doc_roles.register_canonical_role(name, cl)
        self.mapping[str(cl)] = name
        self.app.add_role(name, cl)
        self.writer.app.add_role(name, cl)

    def add_mapping(self, name, cl):
        self.mapping[str(cl)] = name

    def add_config_value(self, name, default, rebuild, types=()):
        if name in self.config.values:
            raise Exception('Config value %r already present' % name)
        if rebuild in (False, True):
            rebuild = rebuild and 'env' or ''
        self.new_options[name] = (default, rebuild, types)
        self.config.values[name] = (default, rebuild, types)

    def get_default_values(self):
        return {k: v[0] for k, v in self.new_options.items()}

    def add_node(self, node, **kwds):
        # type: (nodes.Node, Any) -> None
        if sphinx__display_version__ >= "1.6":
            nodes._add_node_class_names([node.__name__])
            for key, val in kwds.items():
                if not isinstance(val, tuple):
                    continue
                visit, depart = val
                translator = self.writer.app.registry.translators.get(key)
                translators = []
                if translator is not None:
                    translators.append(translator)
                elif key == 'html':
                    from sphinx.writers.html import HTMLTranslator
                    translators.append(HTMLTranslator)
                    if is_html5_writer_available():
                        from sphinx.writers.html5 import HTML5Translator
                        translators.append(HTML5Translator)
                elif key == 'latex':
                    from sphinx.writers.latex import LaTeXTranslator
                    translators.append(LaTeXTranslator)
                elif key == 'text':
                    from sphinx.writers.text import TextTranslator
                    translators.append(TextTranslator)
                elif key == 'man':
                    from sphinx.writers.manpage import ManualPageTranslator
                    translators.append(ManualPageTranslator)
                elif key == 'texinfo':
                    from sphinx.writers.texinfo import TexinfoTranslator
                    translators.append(TexinfoTranslator)

                for translator in translators:
                    setattr(translator, 'visit_' + node.__name__, visit)
                    if depart:
                        setattr(translator, 'depart_' +
                                node.__name__, depart)
        else:
            self.app.add_node(node, **kwds)

    def connect(self, node, func):
        self.mapping_connect[node] = func
        self.app.connect(node, func)
        self.writer.app.connect(node, func)

    def add_domain(self, domain):
        if domain.name in self.domains:
            raise Exception(
                'domain %s already registered' % domain.name)
        self.domains[domain.name] = domain

    def add_event(self, name):
        if name in self._events:
            raise Exception('Event %s already present' % name)
        self._events[name] = ''

    def emit(self, event, *args):
        return self.app.emit(event, *args)

    @staticmethod
    def create(writer="sphinx", directives=None):
        """
        Create a MockApp

        @param      write           ``'sphinx'`` is the only allowed value
        @param      directives      new directives to add (see below)
        @return                     mockapp, writer, list of added nodes

        *directives* is None or a list of 5-uple:

        * a directive name
        * a directive class: see `Sphinx Directive <http://sphinx-doc.org/extdev/tutorial.html>`_, see also @see cl RunPythonDirective as an example
        * a docutils node: see @see cl runpython_node as an example
        * two functions: see @see fn visit_runpython_node, @see fn depart_runpython_node as an example

        """
        if writer not in ("sphinx", "custom"):
            raise NotImplementedError("writer must be 'sphinx' or 'custom'")

        writer = HTMLWriterWithCustomDirectives()
        mockapp = MockSphinxApp(writer, writer.app)

        from sphinx.ext.graphviz import setup as setup_graphviz
        from sphinx.ext.imgmath import setup as setup_math
        from sphinx.ext.todo import setup as setup_todo
        from matplotlib.sphinxext.plot_directive import setup as setup_plot
        from matplotlib.sphinxext.only_directives import setup as setup_only

        # directives from pyquickhelper
        for app in [mockapp, writer.app]:
            setup_blog(app)
            setup_runpython(app)
            setup_sharenet(app)
            setup_todoext(app)
            setup_bigger(app)
            setup_githublink(app)
            setup_runpython(app)
            setup_mathdef(app)
            setup_blocref(app)
            setup_faqref(app)
            setup_exref(app)
            setup_nbref(app)

            # directives from sphinx
            setup_graphviz(app)
            setup_math(app)
            setup_todo(app)

            # don't move this import to the beginning of file
            # it changes matplotlib backend
            setup_plot(app)
            setup_only(app)

        # titles
        title_names = []
        title_names.append("todoext_node")
        title_names.append("todo_node")
        title_names.append("mathdef_node")
        title_names.append("blocref_node")
        title_names.append("faqref_node")
        title_names.append("nbref_node")
        title_names.append("exref_node")

        if directives is not None:
            for tu in directives:
                if len(tu) != 5:
                    raise ValueError(
                        "directives is a list of tuple with 5 elements, check the documentation")
                name, cl, node, f1, f2 = tu
                doc_directives.register_directive(name, cl)
                mockapp.add_directive(name, cl)
                mockapp.add_node(node, html=(f1, f2))
                # not necessary
                # nodes._add_node_class_names([node.__name__])
                writer.connect_directive_node(node.__name__, f1, f2)

        return mockapp, writer, title_names
<commit_msg>documentation<commit_after>"""
@file
@brief Helpers to convert docstring to various format

.. versionadded:: 1.0
"""
from ..sphinxext.sphinx_bigger_extension import setup as setup_bigger
from ..sphinxext.sphinx_githublink_extension import setup as setup_githublink
from ..sphinxext.sphinx_blocref_extension import setup as setup_blocref
from ..sphinxext.sphinx_blog_extension import setup as setup_blog
from ..sphinxext.sphinx_exref_extension import setup as setup_exref
from ..sphinxext.sphinx_faqref_extension import setup as setup_faqref
from ..sphinxext.sphinx_mathdef_extension import setup as setup_mathdef
from ..sphinxext.sphinx_nbref_extension import setup as setup_nbref
from ..sphinxext.sphinx_runpython_extension import setup as setup_runpython
from ..sphinxext.sphinx_sharenet_extension import setup as setup_sharenet
from ..sphinxext.sphinx_todoext_extension import setup as setup_todoext
from .convert_doc_sphinx_helper import HTMLWriterWithCustomDirectives
from docutils import nodes
from docutils.parsers.rst import directives as doc_directives, roles as doc_roles
from sphinx.config import Config
from sphinx import __display_version__ as sphinx__display_version__

try:
    from sphinx.util.docutils import is_html5_writer_available
except ImportError:
    # Available only after Sphinx >= 1.6.1.
    def is_html5_writer_available():
        return False

class MockSphinxApp:
    """
    Mock Sphinx application.
    In memory Sphinx application.
    """

    def __init__(self, writer, app):
        """
        Constructor

        @param      writer      see static method create
        @param      app         see static method create
        """
        self.app = app
        self.new_options = {}
        self.writer = writer
        self.mapping = {"<class 'sphinx.ext.todo.todo_node'>": "todo",
                        "<class 'sphinx.ext.graphviz.graphviz'>": "graphviz",
                        "<class 'sphinx.ext.mathbase.math'>": "math",
                        "<class 'sphinx.ext.mathbase.displaymath'>": "displaymath",
                        "<class 'sphinx.ext.mathbase.eqref'>": "eqref",
                        "<class 'matplotlib.sphinxext.only_directives.html_only'>": "only",
                        "<class 'matplotlib.sphinxext.only_directives.latex_only'>": "only",
                        }
        self.mapping_connect = {}
        self.config = Config(None, None, overrides={}, tags=None)
        self.confdir = "."
        self.doctreedir = "."
        self.srcdir = "."
        self.builder = writer.builder
        self.domains = {}
        self._events = {}

    def add_directive(self, name, cl, *args, **options):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        doc_directives.register_directive(name, cl)
        self.mapping[str(cl)] = name
        self.app.add_directive(name, cl, *args, **options)
        self.writer.app.add_directive(name, cl, *args, **options)

    def add_role(self, name, cl):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        doc_roles.register_canonical_role(name, cl)
        self.mapping[str(cl)] = name
        self.app.add_role(name, cl)
        self.writer.app.add_role(name, cl)

    def add_mapping(self, name, cl):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        self.mapping[str(cl)] = name

    def add_config_value(self, name, default, rebuild, types=()):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        if name in self.config.values:
            raise Exception('Config value %r already present' % name)
        if rebuild in (False, True):
            rebuild = rebuild and 'env' or ''
        self.new_options[name] = (default, rebuild, types)
        self.config.values[name] = (default, rebuild, types)

    def get_default_values(self):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        return {k: v[0] for k, v in self.new_options.items()}

    def add_node(self, node, **kwds):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        # type: (nodes.Node, Any) -> None
        if sphinx__display_version__ >= "1.6":
            nodes._add_node_class_names([node.__name__])
            for key, val in kwds.items():
                if not isinstance(val, tuple):
                    continue
                visit, depart = val
                translator = self.writer.app.registry.translators.get(key)
                translators = []
                if translator is not None:
                    translators.append(translator)
                elif key == 'html':
                    from sphinx.writers.html import HTMLTranslator
                    translators.append(HTMLTranslator)
                    if is_html5_writer_available():
                        from sphinx.writers.html5 import HTML5Translator
                        translators.append(HTML5Translator)
                elif key == 'latex':
                    from sphinx.writers.latex import LaTeXTranslator
                    translators.append(LaTeXTranslator)
                elif key == 'text':
                    from sphinx.writers.text import TextTranslator
                    translators.append(TextTranslator)
                elif key == 'man':
                    from sphinx.writers.manpage import ManualPageTranslator
                    translators.append(ManualPageTranslator)
                elif key == 'texinfo':
                    from sphinx.writers.texinfo import TexinfoTranslator
                    translators.append(TexinfoTranslator)

                for translator in translators:
                    setattr(translator, 'visit_' + node.__name__, visit)
                    if depart:
                        setattr(translator, 'depart_' +
                                node.__name__, depart)
        else:
            self.app.add_node(node, **kwds)

    def connect(self, node, func):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        self.mapping_connect[node] = func
        self.app.connect(node, func)
        self.writer.app.connect(node, func)

    def add_domain(self, domain):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        if domain.name in self.domains:
            raise Exception(
                'domain %s already registered' % domain.name)
        self.domains[domain.name] = domain

    def add_event(self, name):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        if name in self._events:
            raise Exception('Event %s already present' % name)
        self._events[name] = ''

    def emit(self, event, *args):
        """
        See class `Sphinx <https://github.com/sphinx-doc/sphinx/blob/master/sphinx/application.py#L107>`_.
        """
        return self.app.emit(event, *args)

    @staticmethod
    def create(writer="sphinx", directives=None):
        """
        Create a MockApp

        @param      write           ``'sphinx'`` is the only allowed value
        @param      directives      new directives to add (see below)
        @return                     mockapp, writer, list of added nodes

        *directives* is None or a list of 5-uple:

        * a directive name
        * a directive class: see `Sphinx Directive <http://sphinx-doc.org/extdev/tutorial.html>`_,
          see also @see cl RunPythonDirective as an example
        * a docutils node: see @see cl runpython_node as an example
        * two functions: see @see fn visit_runpython_node,
          @see fn depart_runpython_node as an example
        """
        if writer not in ("sphinx", "custom"):
            raise NotImplementedError("writer must be 'sphinx' or 'custom'")

        writer = HTMLWriterWithCustomDirectives()
        mockapp = MockSphinxApp(writer, writer.app)

        from sphinx.ext.graphviz import setup as setup_graphviz
        from sphinx.ext.imgmath import setup as setup_math
        from sphinx.ext.todo import setup as setup_todo
        from matplotlib.sphinxext.plot_directive import setup as setup_plot
        from matplotlib.sphinxext.only_directives import setup as setup_only

        # directives from pyquickhelper
        for app in [mockapp, writer.app]:
            setup_blog(app)
            setup_runpython(app)
            setup_sharenet(app)
            setup_todoext(app)
            setup_bigger(app)
            setup_githublink(app)
            setup_runpython(app)
            setup_mathdef(app)
            setup_blocref(app)
            setup_faqref(app)
            setup_exref(app)
            setup_nbref(app)

            # directives from sphinx
            setup_graphviz(app)
            setup_math(app)
            setup_todo(app)

            # don't move this import to the beginning of file
            # it changes matplotlib backend
            setup_plot(app)
            setup_only(app)

        # titles
        title_names = []
        title_names.append("todoext_node")
        title_names.append("todo_node")
        title_names.append("mathdef_node")
        title_names.append("blocref_node")
        title_names.append("faqref_node")
        title_names.append("nbref_node")
        title_names.append("exref_node")

        if directives is not None:
            for tu in directives:
                if len(tu) != 5:
                    raise ValueError(
                        "directives is a list of tuple with 5 elements, check the documentation")
                name, cl, node, f1, f2 = tu
                doc_directives.register_directive(name, cl)
                mockapp.add_directive(name, cl)
                mockapp.add_node(node, html=(f1, f2))
                # not necessary
                # nodes._add_node_class_names([node.__name__])
                writer.connect_directive_node(node.__name__, f1, f2)

        return mockapp, writer, title_names
<|endoftext|>