half2me / libant

from queue import Queue, Empty

from threading import Lock, Thread, Event

import math

from serial import Serial, SerialException, SerialTimeoutException

import usb

import time

from struct import *

import os

from libAnt.constants import MESSAGE_TX_SYNC, MESSAGE_CHANNEL_BROADCAST_DATA

from libAnt.message import Message, SystemResetMessage

class DriverException(Exception):

    pass

class Logger:

    def __init__(self, logFile: str):

        self._logFile = logFile

        self._log = None

    def __enter__(self):

        self.open()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):

        self.close()

    def open(self):

        if self._log is not None:

            self.close()

        self._log = open(self._logFile, 'wb')

        self.onOpen()

    def close(self):

        if self._log is not None:

            self.beforeClose()

            self._log.close()

            self.afterClose()

    def log(self, data: bytes):

        self._log.write(self.encodeData(data))

    def onOpen(self):

        pass

    def beforeClose(self):

        pass

    def afterClose(self):

        pass

    def encodeData(self, data):

        return data

class PcapLogger(Logger):

    def onOpen(self):

        # write pcap global header

        magic_number = b'\xD4\xC3\xB2\xA1'

        version_major = 2

        version_minor = 4

        thiszone = b'\x00\x00\x00\x00'

        sigfigs = b'\x00\x00\x00\x00'

        snaplen = b'\xFF\x00\x00\x00'

        network = b'\x01\x00\x00\x00'

        pcap_global_header = Struct('<4shh4s4s4s4s')

        self._log.write(

            pcap_global_header.pack(magic_number, version_major, version_minor, thiszone, sigfigs,

                                    snaplen, network))

    def encodeData(self, data):

        timestamp = time.time()

        frac, whole = math.modf(timestamp)

        ts_sec = int(whole).to_bytes(4, byteorder='little')

        ts_usec = int(frac * 1000 * 1000).to_bytes(4, byteorder='little')

        incl_len = len(data)

        orig_len = incl_len

        pcap_packet_header = Struct('<4s4sll').pack(ts_sec, ts_usec, incl_len, orig_len)

        return pcap_packet_header + data

class Driver:

    """

    The driver provides an interface to read and write raw data to and from an ANT+ capable hardware device

    """

    def __init__(self, logger: Logger = None):

        self._lock = Lock()

        self._logger = logger

        self._openTime = None

    def __enter__(self):

        self.open()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):

        self.close()

    def isOpen(self) -> bool:

        with self._lock:

            return self._isOpen()

    def open(self) -> None:

        with self._lock:

            if not self._isOpen():

                self._openTime = time.time()

                if self._logger is not None:

                    self._logger.open()

                self._open()

    def close(self) -> None:

        with self._lock:

            if self._isOpen:

                self._close()

                if self._logger is not None:

                    self._logger.close()

    def reOpen(self) -> None:

        with self._lock:

            if self._isOpen():

                self._close()

            self._open()

    def read(self, timeout=None) -> Message:

        # Splits the string into a list of tokens every n characters

        def splitN(str1, n):

            return [str1[start:start + n] for start in range(0, len(str1), n)]

        if not self.isOpen():

            raise DriverException("Device is closed")

        with self._lock:

            while True:

                sync = self._read(1, timeout=timeout)[0]

                if sync is not MESSAGE_TX_SYNC:

                    continue

                length = self._read(1, timeout=timeout)[0]

                type = self._read(1, timeout=timeout)[0]

                data = self._read(length, timeout=timeout)

                chk = self._read(1, timeout=timeout)[0]

                msg = Message(type, data)

                if self._logger:

                    logMsg = bytearray([sync, length, type])

                    logMsg.extend(data)

                    logMsg.append(chk)

                    self._logger.log(bytes(logMsg))

                if msg.checksum() == chk:

                    return msg

    def write(self, msg: Message) -> None:

        if not self.isOpen():

            raise DriverException("Device is closed")

        with self._lock:

            self._write(msg.encode())

    @abstractmethod

    def _isOpen(self) -> bool:

        pass

    @abstractmethod

    def _open(self) -> None:

        pass

    @abstractmethod

    def _close(self) -> None:

        pass

    @abstractmethod

    def _read(self, count: int, timeout=None) -> bytes:

        pass

    @abstractmethod

    def _write(self, data: bytes) -> None:

        pass

class SerialDriver(Driver):

    """

    An implementation of a serial ANT+ device driver

    """

    def __init__(self, device: str, baudRate: int = 115200, logger: Logger = None):

        super().__init__(logger=logger)

        self._device = device

        self._baudRate = baudRate

        self._serial = None

    def __str__(self):

        if self.isOpen():

            return self._device + " @ " + str(self._baudRate)

        return None

    def _isOpen(self) -> bool:

        return self._serial is None

    def _open(self) -> None:

        try:

            self._serial = Serial(self._device, self._baudRate)

        except SerialException as e:

            raise DriverException(str(e))

        if not self._serial.isOpen():

            raise DriverException("Could not open specified device")

    def _close(self) -> None:

        self._serial.close()

        self._serial = None

    def _read(self, count: int, timeout=None) -> bytes:

        return self._serial.read(count, timeout=timeout)

    def _write(self, data: bytes) -> None:

        try:

            self._serial.write(data)

            self._serial.flush()

        except SerialTimeoutException as e:

            raise DriverException(str(e))

class DummyDriver(Driver):

    def __init__(self, logger: Logger = None):

        super().__init__(logger=logger)

        self._isopen = False

        self._data = Queue()

        msg1 = Message(MESSAGE_CHANNEL_BROADCAST_DATA, b'\x00\x01\x02\x03\x04\x05\x06\x07').encode()

        for b in msg1:

            self._data.put(b)

        msg2 = Message(MESSAGE_CHANNEL_BROADCAST_DATA, b'\x00\xF1\xF2\xF3\xF4\xF5\xF6\xF7').encode()

        for b in msg2:

            self._data.put(b)

        msg3 = Message(MESSAGE_CHANNEL_BROADCAST_DATA, b'\x00\xF9\xFA\xFB\xFC\xFD\xFE\xFF').encode()

        for b in msg3:

            self._data.put(b)

    def _isOpen(self) -> bool:

        return self._isopen

    def _close(self) -> None:

        self._isopen = False

    def _read(self, count: int, timeout=None) -> bytes:

        data = bytearray()

        for i in range(0, count):

            data.append(self._data.get(timeout=timeout))

        return bytes(data)

    def _open(self) -> None:

        self._isopen = True

    def _write(self, data: bytes) -> None:

        pass

class PcapDriver(Driver):

    def __init__(self, pcap, logger: Logger = None):

        super().__init__(logger=logger)

        self._isopen = False

        self._pcap = pcap

        self._buffer = Queue()

        self._loop = None

    class PcapLoop(Thread):

        def __init__(self, pcap, buffer: Queue):

            super().__init__()

            self._stopper = Event()

            self._pcap = pcap

            self._buffer = buffer

        def stop(self) -> None:

            self._stopper.set()

        def run(self) -> None:

            self._pcapfile = open(self._pcap, 'rb')

            self._EOF = os.stat(self._pcap).st_size

            # move file pointer to first packet header

            global_header_length = 24

            self._pcapfile.seek(global_header_length, 0)

            first_ts = 0

            start_time = time.time()

            while not self._stopper.is_set():

                if self._pcapfile.tell() == self._EOF:

                    break

                ts_sec, = unpack('i', self._pcapfile.read(4))

                ts_usec = unpack('i', self._pcapfile.read(4))[0] / 1000000

                if first_ts is 0:

                    first_ts = ts_sec + ts_usec

                ts = ts_sec + ts_usec

                send_time = ts - first_ts

                elapsed_time = time.time() - start_time

                if send_time > (elapsed_time):

                    sleep_time = send_time - elapsed_time

                    time.sleep(sleep_time)

                packet_length = unpack('i', self._pcapfile.read(4))[0]

                self._pcapfile.seek(4, 1)

                for i in range(packet_length):

                    self._buffer.put(self._pcapfile.read(1))

            self._pcapfile.close()

    def _isOpen(self) -> bool:

        return self._isopen

    def _open(self) -> None:

        self._isopen = True

        self._loop = self.PcapLoop(self._pcap, self._buffer)

        self._loop.start()

    def _close(self) -> None:

        self._isopen = False

        if self._loop is not None:

            if self._loop.is_alive():

                self._loop.stop()

                self._loop.join()

        self._loop = None

    def _read(self, count: int, timeout=None) -> bytes:

        result = bytearray()

        while len(result) < count:

            result += self._buffer.get(block=True, timeout=timeout)

        return bytes(result)

    def _write(self, data: bytes) -> None:

        pass