eculocate

eculocate.py (5955B)

---

 #!/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")