Revert "Remove audio path."
[toast/stream2beamer.git] / lagarde.py
1 #!/usr/bin/python3
2 import argparse
3 import asyncio
4 import json
5 import logging
6 import ssl
7 import queue
8 from typing import Optional, List
9
10 import gi
11 import websockets
12
13 gi.require_version('Gst', '1.0')
14 from gi.repository import Gst
15
16 gi.require_version('GstWebRTC', '1.0')
17 from gi.repository import GstWebRTC
18
19 gi.require_version('GstSdp', '1.0')
20 from gi.repository import GstSdp
21
22 gi.require_version('GstRtspServer', '1.0')
23 from gi.repository import Gst, GstRtspServer, GObject, GLib
24
25 log = logging.getLogger(__name__)
26
27
28 class RtspServer:
29     def __init__(self):
30         server = GstRtspServer.RTSPServer()
31         server.set_address("::")
32         server.set_service('8554')  # port as string
33         factory = GstRtspServer.RTSPMediaFactory()
34         # factory.set_launch("intervideosrc ! decodebin ! theoraenc ! queue ! rtptheorapay name=pay0")
35         # factory.set_launch("intervideosrc ! decodebin ! videoconvert ! video/x-raw,format=I420 ! theoraenc ! queue ! rtptheorapay name=pay0")
36         # factory.set_launch("intervideosrc ! decodebin ! videorate ! videoconvert ! video/x-raw,format=I420,framerate=10/1 ! theoraenc ! queue ! rtptheorapay name=pay0")
37         # factory.set_launch("intervideosrc ! decodebin ! videoconvert ! video/x-raw,format=I420 ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
38         factory.set_launch("intervideosrc ! decodebin ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
39         # factory.set_launch("intervideosrc ! decodebin ! videorate ! videoconvert ! video/x-raw,framerate=10/1 ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
40         factory.set_shared(True)
41         mountPoints = server.get_mount_points()
42         mountPoints.add_factory("/cug", factory)
43         server.attach()
44         self.server = server
45
46
47 class Events:
48     def __init__(self):
49         self.sdp_offer = queue.Queue()
50         self.sdp_answer = queue.Queue()
51         self.generated_ice_candidates = queue.Queue()
52         self.received_ice_candidates = queue.Queue()
53         self.sdp_info = queue.Queue()  # (sdp_mids, user_fragments)
54         self.room_left = queue.Queue()
55
56
57 class SignalingClient:
58     def __init__(self, events: Events, uri):
59         self.events = events
60         self.uri = uri
61         self.ssl_context = ssl.SSLContext()
62         self.ssl_context.check_hostname = False
63         self.ssl_context.verify_mode = ssl.CERT_NONE
64         self.session_id = None
65
66     async def receive(self, uri):
67         async for msg in self.websocket:
68             msg_json = json.loads(msg)
69             msg_type = msg_json['Type']
70             msg_value = msg_json['Value']
71             assert self.session_id is None or self.session_id == msg_json['SessionID']
72             if msg_type == 'newSession':
73                 self.session_id = msg_json['SessionID']
74                 log.info(f"New session {self.session_id}")
75             elif msg_type == 'gotOffer':
76                 value_json = json.loads(msg_value)
77                 sdp = value_json['sdp']
78                 log.info(f'Got SDP offer')
79                 log.debug(f'SDP offer:\n{sdp}')
80                 self.events.sdp_offer.put_nowait(sdp)
81             elif msg_type == 'addCallerIceCandidate':
82                 value_json = json.loads(msg_value)
83                 log.info(f'Got ICE candidate')
84                 log.debug(f'ICE candidate: {value_json}')
85                 self.events.received_ice_candidates.put_nowait(value_json)
86             elif msg_type == 'roomNotFound':
87                 log.error(f'The room was not found: {uri}')
88                 return
89             elif msg_type == 'roomClosed':
90                 log.info(f'Oh noes, the room went away (session {self.session_id})!')
91                 self.events.room_left.put_nowait(True)
92                 return
93             else:
94                 log.error(f'Unknown message type {msg_type}')
95
96     async def send(self):
97         sdp_mids = None
98         user_fragments = None
99         while True:
100             if self.events.sdp_answer.qsize() > 0:
101                 sdp_answer = self.events.sdp_answer.get_nowait()
102                 sdp_answer_msg = json.dumps({
103                     'SessionID': self.session_id,
104                     'Type': "gotAnswer",
105                     'Value': json.dumps({
106                         'type': 'answer',
107                         'sdp': sdp_answer
108                     })
109                 })
110                 await self.websocket.send(sdp_answer_msg)
111
112             elif self.events.sdp_info.qsize() > 0:
113                 sdp_mids, user_fragments = self.events.sdp_info.get_nowait()
114
115             elif self.events.generated_ice_candidates.qsize() > 0 \
116                     and sdp_mids is not None and user_fragments is not None:
117                 mlineindex, candidate = self.events.generated_ice_candidates.get_nowait()
118                 sdp_mid = sdp_mids[mlineindex]
119                 user_fragment = user_fragments[mlineindex]
120                 icemsg_value = json.dumps({
121                     "candidate": candidate,
122                     "sdpMid": sdp_mid,
123                     "sdpMLineIndex": mlineindex,
124                     "usernameFragment": user_fragment,
125                 })
126                 icemsg = json.dumps({
127                     'SessionID': self.session_id,
128                     'Type': 'addCalleeIceCandidate',
129                     'Value': icemsg_value,
130                 })
131                 log.info(f'Send ICE candidate')
132                 log.debug(f'ICE candidate: {icemsg_value}')
133                 await self.websocket.send(icemsg)
134
135             else:
136                 await asyncio.sleep(0.2)
137
138     async def run(self):
139         self.session_id = None
140         async with websockets.connect(self.uri, ssl=self.ssl_context, close_timeout=0.5) as self.websocket:
141             receive_task = asyncio.Task(self.receive(self.uri))
142             send_task = asyncio.Task(self.send())
143             done, pending = await asyncio.wait([receive_task, send_task], return_when=asyncio.FIRST_COMPLETED)
144             for task in pending:
145                 task.cancel()
146
147
148 class WebRTCClient:
149     def __init__(self, events: Events):
150         self.events = events
151         self.webrtcbin = Gst.ElementFactory.make('webrtcbin', 'laplace')
152         self.pipe = Gst.Pipeline.new("pipeline")
153         Gst.Bin.do_add_element(self.pipe, self.webrtcbin)
154         self.webrtcbin.connect('on-negotiation-needed', self.on_negotiation_needed)
155         self.webrtcbin.connect('on-ice-candidate', self.on_ice_candidate)
156         self.webrtcbin.connect('pad-added', self.webrtcbin_pad_added)
157
158     def on_negotiation_needed(self, element):
159         log.info('on_negotiation_needed')
160
161     def on_ice_candidate(self, element, mlineindex, candidate):
162         log.info('on_ice_candidate')
163         self.events.generated_ice_candidates.put_nowait((mlineindex, candidate))
164
165     def webrtcbin_pad_added(self, element, pad):
166         log.info('webrtcbin_pad_added')
167         if pad.direction != Gst.PadDirection.SRC:
168             return
169         decodebin = Gst.ElementFactory.make('decodebin')
170         decodebin.connect('pad-added', self.decodebin_pad_added)
171         self.pipe.add(decodebin)
172         decodebin.sync_state_with_parent()
173         self.webrtcbin.link(decodebin)
174
175     def decodebin_pad_added(self, element, pad):
176         log.info('decodebin_pad_added')
177         if not pad.has_current_caps():
178             log.info(pad, 'has no caps, ignoring')
179             return
180
181         caps = pad.get_current_caps()
182         padsize = caps.get_size()
183         for i in range(padsize):
184             s = caps.get_structure(i)  # Gst.Structure
185             name = s.get_name()
186             if name.startswith('video'):
187                 q = Gst.ElementFactory.make('queue')
188                 conv = Gst.ElementFactory.make('videoconvert')
189                 sink = Gst.ElementFactory.make('intervideosink')
190                 self.pipe.add(q)
191                 self.pipe.add(conv)
192                 self.pipe.add(sink)
193                 self.pipe.sync_children_states()
194                 pad.link(q.get_static_pad('sink'))
195                 q.link(conv)
196                 conv.link(sink)
197             elif name.startswith('audio'):
198                 q = Gst.ElementFactory.make('queue')
199                 conv = Gst.ElementFactory.make('audioconvert')
200                 resample = Gst.ElementFactory.make('audioresample')
201                 sink = Gst.ElementFactory.make('autoaudiosink')
202                 self.pipe.add(q)
203                 self.pipe.add(conv)
204                 self.pipe.add(resample)
205                 self.pipe.add(sink)
206                 self.pipe.sync_children_states()
207                 pad.link(q.get_static_pad('sink'))
208                 q.link(conv)
209                 conv.link(resample)
210                 resample.link(sink)
211
212     def set_remote_desciption_done(self, gst_promise):
213         gst_promise = Gst.Promise.new_with_change_func(self.create_answer_done)
214         self.webrtcbin.emit('create-answer', None, gst_promise)
215
216     def create_answer_done(self, gst_promise):
217         reply = gst_promise.get_reply()
218         answer = reply.get_value('answer')
219         sdp_message = answer.sdp
220         mids = [sdp_message.get_media(i).get_attribute_val('mid')
221                      for i in range(sdp_message.medias_len())]
222         user_fragments = [sdp_message.get_media(i).get_attribute_val('ice-ufrag')
223                                for i in range(sdp_message.medias_len())]
224         self.events.sdp_info.put_nowait((mids, user_fragments))
225         sdp_answer = sdp_message.as_text()
226         log.info(f'Send SDP answer')
227         log.debug(f'SDP answer:\n{sdp_answer}')
228         self.events.sdp_answer.put_nowait(sdp_answer)
229         gst_promise = Gst.Promise.new_with_change_func(self.set_local_description_done)
230         self.webrtcbin.emit('set-local-description', answer, gst_promise)
231
232     def set_local_description_done(self, gst_promise):
233         gst_promise.get_reply()
234
235     async def run(self):
236         bus = Gst.Pipeline.get_bus(self.pipe)
237         self.pipe.set_state(Gst.State.PLAYING)
238         try:
239             while True:
240                 if bus.have_pending():
241                     msg = bus.pop()
242                     if msg.type == Gst.MessageType.ERROR:
243                         log.error(f'Error from gstreamer message bus: {msg.get_structure()}')
244                         return
245                     elif msg.type == Gst.MessageType.EOS:  # end of stream
246                         log.info(f'Gstreamer message bus reports end of stream')
247                         return
248                 elif self.events.sdp_offer.qsize() > 0:
249                     sdp_offer = self.events.sdp_offer.get_nowait()
250                     res, sm = GstSdp.SDPMessage.new()
251                     assert res == GstSdp.SDPResult.OK
252                     GstSdp.sdp_message_parse_buffer(bytes(sdp_offer.encode()), sm)
253                     # the three lines above can also be done this way in new versions of GStreamer:
254                     # sm = GstSdp.SDPMessage.new_from_text(sdp_offer)
255                     rd = GstWebRTC.WebRTCSessionDescription.new(GstWebRTC.WebRTCSDPType.OFFER, sm)
256                     gst_promise = Gst.Promise.new_with_change_func(self.set_remote_desciption_done)
257                     self.webrtcbin.emit('set-remote-description', rd, gst_promise)
258
259                 elif self.events.received_ice_candidates.qsize() > 0:
260                     ic = self.events.received_ice_candidates.get_nowait()
261                     if ic['candidate'] != '':
262                         self.webrtcbin.emit('add-ice-candidate', ic['sdpMLineIndex'], ic['candidate'])
263
264                 elif self.events.room_left.qsize() > 0:
265                     self.events.room_left.get_nowait()
266                     return
267
268                 else:
269                     await asyncio.sleep(0.1)
270         finally:
271             self.pipe.set_state(Gst.State.NULL)
272
273
274 async def gstreamer_main_loop():
275     """Does the equivalent of the following lines in an async friendly way:
276         loop = GLib.MainLoop()
277         loop.run()
278     """
279     gst_loop = GLib.MainLoop()
280     context = gst_loop.get_context()
281     while True:
282         events_dispatched = context.iteration(False)
283         await asyncio.sleep(0. if events_dispatched else 0.01)
284
285
286 async def run_repeated(task):
287     while True:
288         await task()
289         await asyncio.sleep(0.1)
290
291
292 async def run(uri):
293     try:
294         events = Events()
295
296         rtsp = RtspServer()
297         webrtc = WebRTCClient(events)
298         signaling = SignalingClient(events, uri)
299
300         main_loop_task = asyncio.Task(gstreamer_main_loop())
301         webrtc_task = asyncio.Task(run_repeated(webrtc.run))
302         signaling_task = asyncio.Task(run_repeated(signaling.run))
303
304         done, pending = await asyncio.wait([main_loop_task, webrtc_task, signaling_task],
305             return_when=asyncio.FIRST_COMPLETED)
306
307         for task in done:
308             task.result()
309         for task in pending:
310             task.cancel()
311     except OSError as e:
312         print(e)
313
314
315 def main():
316     logging.basicConfig(level=logging.INFO, format='%(asctime)-15s %(message)s')
317     parser = argparse.ArgumentParser()
318     parser.add_argument('--uri', default='wss://localhost:1234/ws_connect?id=cug',
319                         help='Signalling server URI')
320     args = parser.parse_args()
321
322     Gst.init(None)
323     asyncio.run(run(args.uri), debug=True)
324
325
326 if __name__ == '__main__':
327     main()