From ff7f4de7c146c742fbea963a2fad33f1bfaf78b2 Mon Sep 17 00:00:00 2001
From: Andrei <b1naryth1ef@gmail.com>
Date: Sun, 25 Sep 2016 07:33:27 -0500
Subject: [PATCH] More voice implementation

---
 disco/voice/client.py    | 132 ++++++++++++++++++++++++++++++++++++---
 examples/basic_plugin.py |   8 ++-
 2 files changed, 130 insertions(+), 10 deletions(-)

diff --git a/disco/voice/client.py b/disco/voice/client.py
index c45c66e..1954ce5 100644
--- a/disco/voice/client.py
+++ b/disco/voice/client.py
@@ -1,4 +1,7 @@
 import gevent
+import socket
+import struct
+import time
 
 from holster.enum import Enum
 from holster.emitter import Emitter
@@ -19,6 +22,11 @@ VoiceState = Enum(
     VOICE_CONNECTED=6,
 )
 
+# TODO:
+#   - player implementation
+#   - encryption
+#   - cleanup
+
 
 class VoiceException(Exception):
     def __init__(self, msg, client):
@@ -26,8 +34,58 @@ class VoiceException(Exception):
         super(VoiceException, self).__init__(msg)
 
 
+class UDPVoiceClient(LoggingClass):
+    def __init__(self, vc):
+        super(UDPVoiceClient, self).__init__()
+        self.vc = vc
+        self.conn = None
+        self.ip = None
+        self.port = None
+        self.run_task = None
+        self.connected = False
+
+    def run(self):
+        while True:
+            self.conn.recvfrom(4096)
+
+    def send(self, data):
+        self.conn.sendto(data, (self.ip, self.port))
+
+    def disconnect(self):
+        self.run_task.kill()
+
+    def connect(self, host, port, timeout=10):
+        self.ip = socket.gethostbyname(host)
+        self.port = port
+
+        self.conn = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+
+        # Send discovery packet
+        packet = bytearray(70)
+        struct.pack_into('>I', packet, 0, self.vc.ssrc)
+        self.send(packet)
+
+        # Wait for a response
+        try:
+            data, addr = gevent.spawn(lambda: self.conn.recvfrom(70)).get(timeout=timeout)
+        except gevent.Timeout:
+            return (None, None)
+
+        # Read IP and port
+        ip = data[4:].split('\x00', 1)[0]
+        port = struct.unpack('<H', data[-2:])[0]
+
+        # Spawn read thread so we don't max buffers
+        self.connected = True
+        self.run_task = gevent.spawn(self.run)
+
+        return (ip, port)
+
+
 class VoiceClient(LoggingClass):
     def __init__(self, channel):
+        super(VoiceClient, self).__init__()
+
         assert(channel.is_voice)
         self.channel = channel
         self.client = self.channel.client
@@ -41,11 +99,25 @@ class VoiceClient(LoggingClass):
         self.connected = gevent.event.Event()
         self.token = None
         self.endpoint = None
+        self.ssrc = None
+        self.port = None
 
         self.update_listener = None
 
         # Websocket connection
         self.ws = None
+        self.heartbeat_task = None
+
+    def heartbeat(self, interval):
+        while True:
+            self.send(VoiceOPCode.HEARTBEAT, time.time() * 1000)
+            gevent.sleep(interval / 1000)
+
+    def set_speaking(self, value):
+        self.send(VoiceOPCode.SPEAKING, {
+            'speaking': value,
+            'delay': 0,
+        })
 
     def send(self, op, data):
         self.ws.send(dumps({
@@ -54,10 +126,36 @@ class VoiceClient(LoggingClass):
         }))
 
     def on_voice_ready(self, data):
-        print data
+        self.state = VoiceState.CONNECTING
+        self.ssrc = data['ssrc']
+        self.port = data['port']
+
+        self.heartbeat_task = gevent.spawn(self.heartbeat, data['heartbeat_interval'])
+
+        self.udp = UDPVoiceClient(self)
+        ip, port = self.udp.connect(self.endpoint, self.port)
+
+        if not ip:
+            self.disconnect()
+            return
+
+        self.send(VoiceOPCode.SELECT_PROTOCOL, {
+            'protocol': 'udp',
+            'data': {
+                'port': port,
+                'address': ip,
+                'mode': 'plain'
+            }
+        })
 
     def on_voice_sdp(self, data):
-        print data
+        # Toggle speaking state so clients learn of our SSRC
+        self.set_speaking(True)
+        self.set_speaking(False)
+        gevent.sleep(0.25)
+
+        self.state = VoiceState.CONNECTED
+        self.connected.set()
 
     def on_voice_server_update(self, data):
         if self.channel.guild_id != data.guild_id or not data.token:
@@ -69,9 +167,9 @@ class VoiceClient(LoggingClass):
         self.token = data.token
         self.state = VoiceState.AUTHENTICATING
 
-        self.endpoint = 'wss://{}'.format(data.endpoint.split(':', 1)[0])
+        self.endpoint = data.endpoint.split(':', 1)[0]
         self.ws = Websocket(
-            self.endpoint,
+            'wss://' + self.endpoint,
             on_message=self.on_message,
             on_error=self.on_error,
             on_open=self.on_open,
@@ -85,7 +183,7 @@ class VoiceClient(LoggingClass):
         except:
             self.log.exception('Failed to parse voice gateway message: ')
 
-        self.packets.emit(VoiceOPCode[data['op']], data)
+        self.packets.emit(VoiceOPCode[data['op']], data['d'])
 
     def on_error(self, ws, err):
         # TODO
@@ -99,9 +197,9 @@ class VoiceClient(LoggingClass):
             'token': self.token
         })
 
-    def on_close(self, ws):
+    def on_close(self, ws, code, error):
         # TODO
-        self.log.warning('Voice websocket disconnected')
+        self.log.warning('Voice websocket disconnected (%s, %s)', code, error)
 
     def connect(self, timeout=5, mute=False, deaf=False):
         self.state = VoiceState.AWAITING_ENDPOINT
@@ -117,3 +215,23 @@ class VoiceClient(LoggingClass):
 
         if not self.connected.wait(timeout) or self.state != VoiceState.CONNECTED:
             raise VoiceException('Failed to connect to voice', self)
+
+    def disconnect(self):
+        self.state = VoiceState.DISCONNECTED
+
+        if self.heartbeat_task:
+            self.heartbeat_task.kill()
+            self.heartbeat_task = None
+
+        if self.ws and self.ws.sock.connected:
+            self.ws.close()
+
+        if self.udp and self.udp.connected:
+            self.udp.disconnect()
+
+        self.client.gw.send(OPCode.VOICE_STATE_UPDATE, {
+            'self_mute': False,
+            'self_deaf': False,
+            'guild_id': int(self.channel.guild_id),
+            'channel_id': None,
+        })
diff --git a/examples/basic_plugin.py b/examples/basic_plugin.py
index 4bfd16f..0708a57 100644
--- a/examples/basic_plugin.py
+++ b/examples/basic_plugin.py
@@ -1,3 +1,5 @@
+import gevent
+
 from disco.cli import disco_main
 from disco.bot import Bot
 from disco.bot.plugin import Plugin
@@ -56,9 +58,9 @@ class BasicPlugin(Plugin):
             event.msg.reply('You are not connected to voice')
             return
 
-        print vs.channel
-        print vs.channel_id
-        print vs.channel.connect()
+        vc = vs.channel.connect()
+        gevent.sleep(1)
+        vc.disconnect()
 
 if __name__ == '__main__':
     bot = Bot(disco_main())