| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  | #!/usr/bin/env python | 
					
						
							|  |  |  | # Tool to subscribe to motion data and log it to a disk file | 
					
						
							|  |  |  | # | 
					
						
							|  |  |  | # Copyright (C) 2020-2021  Kevin O'Connor <kevin@koconnor.net> | 
					
						
							|  |  |  | # | 
					
						
							|  |  |  | # This file may be distributed under the terms of the GNU GPLv3 license. | 
					
						
							|  |  |  | import sys, os, optparse, socket, select, json, errno, time, zlib | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | INDEX_UPDATE_TIME = 5.0 | 
					
						
							|  |  |  | ClientInfo = {'program': 'motan_data_logger', 'version': 'v0.1'} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | def webhook_socket_create(uds_filename): | 
					
						
							|  |  |  |     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) | 
					
						
							|  |  |  |     sock.setblocking(0) | 
					
						
							|  |  |  |     sys.stderr.write("Waiting for connect to %s\n" % (uds_filename,)) | 
					
						
							|  |  |  |     while 1: | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             sock.connect(uds_filename) | 
					
						
							|  |  |  |         except socket.error as e: | 
					
						
							|  |  |  |             if e.errno == errno.ECONNREFUSED: | 
					
						
							|  |  |  |                 time.sleep(0.1) | 
					
						
							|  |  |  |                 continue | 
					
						
							|  |  |  |             sys.stderr.write("Unable to connect socket %s [%d,%s]\n" | 
					
						
							|  |  |  |                              % (uds_filename, e.errno, | 
					
						
							|  |  |  |                                 errno.errorcode[e.errno])) | 
					
						
							|  |  |  |             sys.exit(-1) | 
					
						
							|  |  |  |         break | 
					
						
							|  |  |  |     sys.stderr.write("Connection.\n") | 
					
						
							|  |  |  |     return sock | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class LogWriter: | 
					
						
							|  |  |  |     def __init__(self, filename): | 
					
						
							|  |  |  |         self.file = open(filename, "wb") | 
					
						
							|  |  |  |         self.comp = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, | 
					
						
							|  |  |  |                                      zlib.DEFLATED, 31) | 
					
						
							|  |  |  |         self.raw_pos = self.file_pos = 0 | 
					
						
							|  |  |  |     def add_data(self, data): | 
					
						
							|  |  |  |         d = self.comp.compress(data + b"\x03") | 
					
						
							|  |  |  |         self.file.write(d) | 
					
						
							|  |  |  |         self.file_pos += len(d) | 
					
						
							|  |  |  |         self.raw_pos += len(data) + 1 | 
					
						
							|  |  |  |     def flush(self, flag=zlib.Z_FULL_FLUSH): | 
					
						
							|  |  |  |         if not self.raw_pos: | 
					
						
							|  |  |  |             return self.file_pos | 
					
						
							|  |  |  |         d = self.comp.flush(flag) | 
					
						
							|  |  |  |         self.file.write(d) | 
					
						
							|  |  |  |         self.file_pos += len(d) | 
					
						
							|  |  |  |         return self.file_pos | 
					
						
							|  |  |  |     def close(self): | 
					
						
							|  |  |  |         self.flush(zlib.Z_FINISH) | 
					
						
							|  |  |  |         self.file.close() | 
					
						
							|  |  |  |         self.file = None | 
					
						
							|  |  |  |         self.comp = None | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class DataLogger: | 
					
						
							|  |  |  |     def __init__(self, uds_filename, log_prefix): | 
					
						
							|  |  |  |         # IO | 
					
						
							|  |  |  |         self.webhook_socket = webhook_socket_create(uds_filename) | 
					
						
							|  |  |  |         self.poll = select.poll() | 
					
						
							|  |  |  |         self.poll.register(self.webhook_socket, select.POLLIN | select.POLLHUP) | 
					
						
							|  |  |  |         self.socket_data = b"" | 
					
						
							|  |  |  |         # Data log | 
					
						
							|  |  |  |         self.logger = LogWriter(log_prefix + ".json.gz") | 
					
						
							|  |  |  |         self.index = LogWriter(log_prefix + ".index.gz") | 
					
						
							|  |  |  |         # Handlers | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |         self.query_handlers = {} | 
					
						
							|  |  |  |         self.async_handlers = {} | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |         # get_status databasing | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |         self.db = {} | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |         self.next_index_time = 0. | 
					
						
							|  |  |  |         # Start login process | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |         self.send_query("info", "info", {"client_info": ClientInfo}, | 
					
						
							|  |  |  |                         self.handle_info) | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |     def error(self, msg): | 
					
						
							|  |  |  |         sys.stderr.write(msg + "\n") | 
					
						
							|  |  |  |     def finish(self, msg): | 
					
						
							|  |  |  |         self.error(msg) | 
					
						
							|  |  |  |         self.logger.close() | 
					
						
							|  |  |  |         self.index.close() | 
					
						
							|  |  |  |         sys.exit(0) | 
					
						
							|  |  |  |     # Unix Domain Socket IO | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |     def send_query(self, msg_id, method, params, cb): | 
					
						
							|  |  |  |         self.query_handlers[msg_id] = cb | 
					
						
							|  |  |  |         msg = {"id": msg_id, "method": method, "params": params} | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |         cm = json.dumps(msg, separators=(',', ':')).encode() | 
					
						
							|  |  |  |         self.webhook_socket.send(cm + b"\x03") | 
					
						
							|  |  |  |     def process_socket(self): | 
					
						
							|  |  |  |         data = self.webhook_socket.recv(4096) | 
					
						
							|  |  |  |         if not data: | 
					
						
							|  |  |  |             self.finish("Socket closed") | 
					
						
							|  |  |  |         parts = data.split(b"\x03") | 
					
						
							|  |  |  |         parts[0] = self.socket_data + parts[0] | 
					
						
							|  |  |  |         self.socket_data = parts.pop() | 
					
						
							|  |  |  |         for part in parts: | 
					
						
							|  |  |  |             try: | 
					
						
							|  |  |  |                 msg = json.loads(part) | 
					
						
							|  |  |  |             except: | 
					
						
							|  |  |  |                 self.error("ERROR: Unable to parse line") | 
					
						
							|  |  |  |                 continue | 
					
						
							|  |  |  |             self.logger.add_data(part) | 
					
						
							|  |  |  |             msg_q = msg.get("q") | 
					
						
							|  |  |  |             if msg_q is not None: | 
					
						
							|  |  |  |                 hdl = self.async_handlers.get(msg_q) | 
					
						
							|  |  |  |                 if hdl is not None: | 
					
						
							|  |  |  |                     hdl(msg, part) | 
					
						
							|  |  |  |                 continue | 
					
						
							|  |  |  |             msg_id = msg.get("id") | 
					
						
							|  |  |  |             hdl = self.query_handlers.get(msg_id) | 
					
						
							|  |  |  |             if hdl is not None: | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |                 del self.query_handlers[msg_id] | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |                 hdl(msg, part) | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |                 if not self.query_handlers: | 
					
						
							|  |  |  |                     self.flush_index() | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |                 continue | 
					
						
							|  |  |  |             self.error("ERROR: Message with unknown id") | 
					
						
							|  |  |  |     def run(self): | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             while 1: | 
					
						
							|  |  |  |                 res = self.poll.poll(1000.) | 
					
						
							|  |  |  |                 for fd, event in res: | 
					
						
							|  |  |  |                     if fd == self.webhook_socket.fileno(): | 
					
						
							|  |  |  |                         self.process_socket() | 
					
						
							|  |  |  |         except KeyboardInterrupt as e: | 
					
						
							|  |  |  |             self.finish("Keyboard Interrupt") | 
					
						
							|  |  |  |     # Query response handlers | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |     def send_subscribe(self, msg_id, method, params, cb=None, async_cb=None): | 
					
						
							|  |  |  |         if cb is None: | 
					
						
							|  |  |  |             cb = self.handle_dump | 
					
						
							|  |  |  |         if async_cb is not None: | 
					
						
							|  |  |  |             self.async_handlers[msg_id] = async_cb | 
					
						
							|  |  |  |         params["response_template"] = {"q": msg_id} | 
					
						
							|  |  |  |         self.send_query(msg_id, method, params, cb) | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |     def handle_info(self, msg, raw_msg): | 
					
						
							|  |  |  |         if msg["result"]["state"] != "ready": | 
					
						
							|  |  |  |             self.finish("Klipper not in ready state") | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |         self.send_query("list", "objects/list", {}, self.handle_list) | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |     def handle_list(self, msg, raw_msg): | 
					
						
							|  |  |  |         subreq = {o: None for o in msg["result"]["objects"]} | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |         self.send_subscribe("status", "objects/subscribe", {"objects": subreq}, | 
					
						
							|  |  |  |                             self.handle_subscribe, self.handle_async_db) | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |     def handle_subscribe(self, msg, raw_msg): | 
					
						
							|  |  |  |         result = msg["result"] | 
					
						
							|  |  |  |         self.next_index_time = result["eventtime"] + INDEX_UPDATE_TIME | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |         self.db["status"] = status = result["status"] | 
					
						
							|  |  |  |         # Subscribe to trapq and stepper queue updates | 
					
						
							|  |  |  |         motion_report = status.get("motion_report", {}) | 
					
						
							|  |  |  |         for trapq in motion_report.get("trapq", []): | 
					
						
							|  |  |  |             self.send_subscribe("trapq:" + trapq, "motion_report/dump_trapq", | 
					
						
							|  |  |  |                                 {"name": trapq}) | 
					
						
							|  |  |  |         for stepper in motion_report.get("steppers", []): | 
					
						
							|  |  |  |             self.send_subscribe("stepq:" + stepper, | 
					
						
							|  |  |  |                                 "motion_report/dump_stepper", {"name": stepper}) | 
					
						
							| 
									
										
										
										
											2021-08-12 15:59:41 -04:00
										 |  |  |         # Subscribe to additional sensor data | 
					
						
							| 
									
										
										
										
											2023-12-15 18:04:17 -05:00
										 |  |  |         stypes = ["adxl345", "lis2dw", "mpu9250", "angle"] | 
					
						
							| 
									
										
										
										
											2023-12-15 19:05:25 -05:00
										 |  |  |         stypes = {st:st for st in stypes} | 
					
						
							|  |  |  |         stypes['probe_eddy_current'] = 'ldc1612' | 
					
						
							| 
									
										
										
										
											2021-08-12 15:59:41 -04:00
										 |  |  |         config = status["configfile"]["settings"] | 
					
						
							|  |  |  |         for cfgname in config.keys(): | 
					
						
							| 
									
										
										
										
											2023-12-15 19:05:25 -05:00
										 |  |  |             for capprefix, st in sorted(stypes.items()): | 
					
						
							|  |  |  |                 if cfgname == capprefix or cfgname.startswith(capprefix + " "): | 
					
						
							| 
									
										
										
										
											2023-12-15 18:04:17 -05:00
										 |  |  |                     aname = cfgname.split()[-1] | 
					
						
							|  |  |  |                     lname = "%s:%s" % (st, aname) | 
					
						
							|  |  |  |                     qcmd = "%s/dump_%s" % (st, st) | 
					
						
							|  |  |  |                     self.send_subscribe(lname, qcmd, {"sensor": aname}) | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |     def handle_dump(self, msg, raw_msg): | 
					
						
							|  |  |  |         msg_id = msg["id"] | 
					
						
							| 
									
										
										
										
											2021-09-26 12:24:40 -04:00
										 |  |  |         if "result" not in msg: | 
					
						
							|  |  |  |             self.error("Unable to subscribe to '%s': %s" | 
					
						
							|  |  |  |                        % (msg_id, msg.get("error", {}).get("message", ""))) | 
					
						
							|  |  |  |             return | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |         self.db.setdefault("subscriptions", {})[msg_id] = msg["result"] | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |     def flush_index(self): | 
					
						
							|  |  |  |         self.db['file_position'] = self.logger.flush() | 
					
						
							|  |  |  |         self.index.add_data(json.dumps(self.db, separators=(',', ':')).encode()) | 
					
						
							|  |  |  |         self.db = {"status": {}} | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  |     def handle_async_db(self, msg, raw_msg): | 
					
						
							|  |  |  |         params = msg["params"] | 
					
						
							|  |  |  |         db_status = self.db['status'] | 
					
						
							|  |  |  |         for k, v in params.get("status", {}).items(): | 
					
						
							|  |  |  |             db_status.setdefault(k, {}).update(v) | 
					
						
							|  |  |  |         eventtime = params['eventtime'] | 
					
						
							| 
									
										
										
										
											2021-08-23 11:01:56 -04:00
										 |  |  |         if eventtime >= self.next_index_time: | 
					
						
							|  |  |  |             self.next_index_time = eventtime + INDEX_UPDATE_TIME | 
					
						
							|  |  |  |             self.flush_index() | 
					
						
							| 
									
										
										
										
											2021-07-23 23:00:02 -04:00
										 |  |  | 
 | 
					
						
							|  |  |  | def nice(): | 
					
						
							|  |  |  |     try: | 
					
						
							|  |  |  |         # Try to re-nice writing process | 
					
						
							|  |  |  |         os.nice(10) | 
					
						
							|  |  |  |     except: | 
					
						
							|  |  |  |         pass | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | def main(): | 
					
						
							|  |  |  |     usage = "%prog [options] <socket filename> <log name>" | 
					
						
							|  |  |  |     opts = optparse.OptionParser(usage) | 
					
						
							|  |  |  |     options, args = opts.parse_args() | 
					
						
							|  |  |  |     if len(args) != 2: | 
					
						
							|  |  |  |         opts.error("Incorrect number of arguments") | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |     nice() | 
					
						
							|  |  |  |     dl = DataLogger(args[0], args[1]) | 
					
						
							|  |  |  |     dl.run() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | if __name__ == '__main__': | 
					
						
							|  |  |  |     main() |