eculocate/client/eculocate.py

#!/usr/bin/env python3

import time
import sys
import socket
import threading
import select
import datetime
from pathlib import Path
import os
import struct
import secrets
import nacl
import nacl.hash
from nacl.encoding import RawEncoder
import configparser

from nacl.signing import SigningKey

from zeroconf import ServiceBrowser, ServiceListener, Zeroconf

class ZListener(ServiceListener):
    def __init__(self, client):
        self.client = client;

    def update_service(self, zc: Zeroconf, type_: str, name: str) -> None:
        print(f"Service {name} updated")
        self.client.on_services_changed()

    def remove_service(self, zc: Zeroconf, type_: str, name: str) -> None:
        print(f"Service {name} removed")
        self.client.on_services_changed()

    def add_service(self, zc: Zeroconf, type_: str, name: str) -> None:
        info = zc.get_service_info(type_, name)
        print(f"Service {name} added, service info: {info}")
        self.client.host_addresses[name] = (socket.inet_ntoa(info.addresses[0]), info.port)
        self.client.on_services_changed()

class Config():
    def __init__(self, pathname):
        config = configparser.ConfigParser(allow_unnamed_section=True)
        config.read(pathname)
        self.config = config
        self._ota_private_key = None

    def ble_private_key(self):
        # FIXME unimplemented
        pathname = self.config.get(configparser.UNNAMED_SECTION, 'BlePrivateKeyFile')

    def ble_peer_key(self):
        # FIXME unimplemented
        pathname = self.config.get(configparser.UNNAMED_SECTION, 'BlePeerKeyFile')

    def ota_private_key(self):
        if not self._ota_private_key:
            pathname = self.config.get(configparser.UNNAMED_SECTION, 'OtaPrivateKeyFile')
            with open(pathname, "rb") as f:
                self._ota_private_key = SigningKey(f.read(), encoder=RawEncoder)

        return self._ota_private_key

    def output_path(self):
        pathname = Path(self.config.get(configparser.UNNAMED_SECTION, 'OutputPath'))
        pathname.mkdir(0o777, False, True)
        return pathname

class Client():
    def __init__(self, config):
        self.selected_service_name = None
        self.bytes_recorded = 0
        self.private_key = config.ota_private_key()
        name = f"{int(time.time())}.raw"
        self._output_pathname = config.output_path() / name

        self.started = False
        self.host_addresses = {}
        zeroconf = Zeroconf()
        ServiceBrowser(zeroconf, "_keihin._udp.local.", ZListener(self))

    def knows_service(self, servicename):
        return not (not (self.host_addresses.get(servicename, False)))

    def output_pathname(self):
        return self._output_pathname

    def recording_thread_main(self):
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        addr = self.host_addresses[self.selected_service_name]
        print(self.selected_service_name)
        print(addr)
        sock.connect(addr)
        sock.setblocking(0)
        interval = 500

        self.bytes_recorded = 0

        with open(self.output_pathname(), "ab") as outf:
            while self.started:
                print("sending subscribe")
                b = self.session_key + bytes([interval >> 8, interval & 0xff, 0x54, 0x11, 0, 25] )
                sock.send(b)
                timeout = time.time() + 60
                while self.started and time.time() < timeout:
                    ready = select.select([sock], [], [], timeout - time.time())
                    if ready[0]:
                        data = sock.recv(4000)
                        self.bytes_recorded = self.bytes_recorded + len(data)
                        outf.write(data)
                        outf.flush()
        sock.close()
        print("recording thread terminated")

    def select_service_name(self, name):
        if name == None:
            self.started = False
        self.selected_service_name = name
        self.on_services_changed()

    def register_session(self):
        session_key = secrets.token_bytes(32)
        sock = socket.create_connection(self.host_addresses[self.selected_service_name])
        sig = self.private_key.sign(session_key, encoder=RawEncoder)

        sock.send(b"KEY0" + session_key + sig.signature)
        response = sock.recv(200)
        print(response.decode(errors="replace"))
        sock.close()
        return session_key

    def start_recording(self):
        if self.selected_service_name != None:
            self.session_key = self.register_session()
            self.started = True
            print("starting record")
            peer_thread = threading.Thread(target=self.recording_thread_main, daemon = True)
            peer_thread.start()
            return True
        else:
            return False

    def stop_recording(self):
        print("stopping record")
        self.started = False

    def on_services_changed(self):
        # override this in subclass if you want to know when
        # service instances appear/disappear
        return None


if __name__ == '__main__':
    import pprint

    exists = False

    conf_file = sys.argv[1]
    name = sys.argv[2]
    timeout = int(sys.argv[3])

    fqname = f"{name}._keihin._udp.local."

    class MyClient(Client):
        def on_services_changed(self):
            print("services changed: now known ")
            pprint.pp(self.host_addresses)
            if self.host_addresses.get(fqname, False):
                print(f"found {fqname}")

    config = Config(conf_file)
    client = MyClient(config)

    while not exists:
        time.sleep(2)
        exists = client.knows_service(fqname)
    client.select_service_name(fqname)

    print("record")
    start_time = time.time()
    client.start_recording()
    while time.time() < start_time + timeout:
        time.sleep(1)
        print(f"{client.bytes_recorded} bytes received")

    client.stop_recording()
    print("fin")