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