record audio from snips.ai mqtt traffic

recorder.py

# record audio from snips.ai mqtt traffic

import paho.mqtt.client as mqtt
import json
import struct
import wave
import datetime
import os

VC_SERVER = 'localhost'
VC_PORT = 1883

record_running = False
record = wave.Wave_write

def on_connect(client, userdata, flags, rc):
	print('Connected to VC system')
	mqtt.subscribe('hermes/audioServer/default/audioFrame')
	#sc_say("recording")

def on_message(client, userdata, msg):
	if msg.topic == "hermes/audioServer/default/audioFrame":
		start_record(msg)

def start_record(msg):
	global record_running
	global record

	riff, size, fformat = struct.unpack('<4sI4s', msg.payload[:12])
	if riff != b'RIFF':
		print("RIFF parse error")
		return
	if fformat != b'WAVE':
		print("FORMAT parse error")
		return
	#print("size: %d" % size)

	# Data Header
	chunk_header = msg.payload[12:20]
	subchunkid, subchunksize = struct.unpack('<4sI', chunk_header)
	
	if (subchunkid == b'fmt '):
		aformat, channels, samplerate, byterate, blockalign, bps = struct.unpack('HHIIHH', msg.payload[20:36])
		bitrate = (samplerate * channels * bps) / 1024
		#print("Format: %i, Channels %i, Sample Rate: %i, Kbps: %i" % (aformat, channels, samplerate, bitrate))

	if(not record_running):
		record_running = True
		record = wave.open(os.path.join("recordings" , datetime.datetime.now().strftime("%Y%m%d-%H%M%S")+".wav"), "wb")
		record.setnchannels(channels)
		record.setframerate(samplerate)
		record.setsampwidth(2)

	chunkOffset = 36
	while (chunkOffset < size):
		subchunk2id, subchunk2size = struct.unpack('<4sI', msg.payload[chunkOffset:chunkOffset+8])
		chunkOffset += 8
		#print("chunk id: %s, size: %i" % (subchunk2id, subchunk2size))
		if (subchunk2id == b'data'):
			if(record_running):
				record.writeframes(msg.payload[chunkOffset:chunkOffset+subchunk2size])
		else:
			print("Data: %s" % msg.payload[chunkOffset:chunkOffset+subchunk2size])

		chunkOffset = chunkOffset + subchunk2size + 8

def stop_recording():
	global record_running
	global record
	record_running = False
	record.close()

def main():
	global record_running
	mqtt.on_connect = on_connect
	mqtt.on_message = on_message
	mqtt.connect(VC_SERVER, VC_PORT)

	mqtt.loop_forever()

if __name__ == "__main__":
	mqtt = mqtt.Client()
	main()