Fix indentation.
[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, rtmp_uri: str):
127         self.events = events
128         self.rtmp_uri = rtmp_uri
129         self.webrtcbin = Gst.ElementFactory.make('webrtcbin', 'laplace')
130         self.pipe = Gst.Pipeline.new("pipeline")
131         Gst.Bin.do_add_element(self.pipe, self.webrtcbin)
132         self.webrtcbin.connect('on-negotiation-needed', self.on_negotiation_needed)
133         self.webrtcbin.connect('on-ice-candidate', self.on_ice_candidate)
134         self.webrtcbin.connect('pad-added', self.webrtcbin_pad_added)
135
136     def on_negotiation_needed(self, element):
137         log.info('on_negotiation_needed')
138
139     def on_ice_candidate(self, element, mlineindex, candidate):
140         log.info('on_ice_candidate')
141         self.events.generated_ice_candidates.put_nowait((mlineindex, candidate))
142
143     def webrtcbin_pad_added(self, element, pad):
144         log.info('webrtcbin_pad_added')
145         if pad.direction != Gst.PadDirection.SRC:
146             return
147         decodebin = Gst.ElementFactory.make('decodebin')
148         decodebin.connect('pad-added', self.decodebin_pad_added)
149         self.pipe.add(decodebin)
150         decodebin.sync_state_with_parent()
151         self.webrtcbin.link(decodebin)
152
153     def decodebin_pad_added(self, element, pad):
154         log.info('decodebin_pad_added')
155         if not pad.has_current_caps():
156             log.info(pad, 'has no caps, ignoring')
157             return
158         caps = pad.get_current_caps()
159         padsize = caps.get_size()
160
161         for i in range(padsize):
162             s = caps.get_structure(i)  # Gst.Structure
163             name = s.get_name()
164             if name.startswith('video'):
165                 q = Gst.ElementFactory.make('queue')
166                 conv = Gst.ElementFactory.make('videoconvert')
167                 enc = Gst.ElementFactory.make('x264enc')
168                 enc.set_property('bitrate', 1000)
169                 enc.set_property('tune', 'zerolatency')
170                 capsfilter = Gst.ElementFactory.make('capsfilter')
171                 capsfilter.set_properties(Gst.Caps.from_string('video/x-h264,stream-format=(string)avc'))
172                 flmux = Gst.ElementFactory.make('flvmux')
173                 sink = Gst.ElementFactory.make('rtmpsink')
174                 sink.set_property('location', self.rtmp_uri)
175                 assert q and conv and enc and capsfilter and flmux and sink
176
177                 self.pipe.add(q)
178                 self.pipe.add(conv)
179                 self.pipe.add(enc)
180                 self.pipe.add(capsfilter)
181                 self.pipe.add(flmux)
182                 self.pipe.add(sink)
183
184                 q_pad_sink = q.get_static_pad('sink')
185                 assert q_pad_sink
186                 pad_link_return = pad.link(q_pad_sink)
187                 assert pad_link_return == Gst.PadLinkReturn.OK
188
189                 ok = q.link(conv)
190                 assert ok
191                 ok = conv.link(enc)
192                 assert ok
193                 ok = enc.link(capsfilter)
194                 assert ok
195                 ok = capsfilter.link(flmux)
196                 assert ok
197                 ok = flmux.link(sink)
198                 assert ok
199                 self.pipe.set_state(Gst.State.PLAYING)
200                 self.pipe.sync_children_states()
201
202             elif name.startswith('audio'):
203                 q = Gst.ElementFactory.make('queue')
204                 conv = Gst.ElementFactory.make('audioconvert')
205                 resample = Gst.ElementFactory.make('audioresample')
206                 sink = Gst.ElementFactory.make('autoaudiosink')
207                 self.pipe.add(q)
208                 self.pipe.add(conv)
209                 self.pipe.add(resample)
210                 self.pipe.add(sink)
211                 self.pipe.sync_children_states()
212                 pad.link(q.get_static_pad('sink'))
213                 q.link(conv)
214                 conv.link(resample)
215                 resample.link(sink)
216
217     def set_remote_desciption_done(self, gst_promise):
218         gst_promise = Gst.Promise.new_with_change_func(self.create_answer_done)
219         self.webrtcbin.emit('create-answer', None, gst_promise)
220
221     def create_answer_done(self, gst_promise):
222         reply = gst_promise.get_reply()
223         answer = reply.get_value('answer')
224         sdp_message = answer.sdp
225         mids = [sdp_message.get_media(i).get_attribute_val('mid')
226                 for i in range(sdp_message.medias_len())]
227         user_fragments = [sdp_message.get_media(i).get_attribute_val('ice-ufrag')
228                           for i in range(sdp_message.medias_len())]
229         self.events.sdp_info.put_nowait((mids, user_fragments))
230         sdp_answer = sdp_message.as_text()
231         log.info(f'Send SDP answer')
232         log.debug(f'SDP answer:\n{sdp_answer}')
233         self.events.sdp_answer.put_nowait(sdp_answer)
234         gst_promise = Gst.Promise.new_with_change_func(self.set_local_description_done)
235         self.webrtcbin.emit('set-local-description', answer, gst_promise)
236
237     def set_local_description_done(self, gst_promise):
238         gst_promise.get_reply()
239
240     async def run(self):
241         bus = Gst.Pipeline.get_bus(self.pipe)
242         self.pipe.set_state(Gst.State.PLAYING)
243         try:
244             while True:
245                 if bus.have_pending():
246                     msg = bus.pop()
247                     if msg.type == Gst.MessageType.ERROR:
248                         log.error(f'Error from gstreamer message bus: {msg.get_structure()}')
249                         return
250                     elif msg.type == Gst.MessageType.EOS:  # end of stream
251                         log.info(f'Gstreamer message bus reports end of stream')
252                         return
253                 elif self.events.sdp_offer.qsize() > 0:
254                     sdp_offer = self.events.sdp_offer.get_nowait()
255                     res, sm = GstSdp.SDPMessage.new()
256                     assert res == GstSdp.SDPResult.OK
257                     GstSdp.sdp_message_parse_buffer(bytes(sdp_offer.encode()), sm)
258                     # the three lines above can also be done this way in new versions of GStreamer:
259                     # sm = GstSdp.SDPMessage.new_from_text(sdp_offer)
260                     rd = GstWebRTC.WebRTCSessionDescription.new(GstWebRTC.WebRTCSDPType.OFFER, sm)
261                     gst_promise = Gst.Promise.new_with_change_func(self.set_remote_desciption_done)
262                     self.webrtcbin.emit('set-remote-description', rd, gst_promise)
263
264                 elif self.events.received_ice_candidates.qsize() > 0:
265                     ic = self.events.received_ice_candidates.get_nowait()
266                     if ic['candidate'] != '':
267                         self.webrtcbin.emit('add-ice-candidate', ic['sdpMLineIndex'], ic['candidate'])
268
269                 elif self.events.room_left.qsize() > 0:
270                     self.events.room_left.get_nowait()
271                     return
272
273                 else:
274                     await asyncio.sleep(0.1)
275         finally:
276             self.pipe.set_state(Gst.State.NULL)
277
278
279 async def run_repeated(task):
280     while True:
281         await task()
282         await asyncio.sleep(0.1)
283
284
285 async def run(laplace_uri: str, rtmp_uri: str):
286     try:
287         events = Events()
288         webrtc = WebRTCClient(events, rtmp_uri)
289         signaling = SignalingClient(events, laplace_uri)
290
291         webrtc_task = asyncio.Task(webrtc.run())
292         signaling_task = asyncio.Task(signaling.run())
293
294         done, pending = await asyncio.wait([webrtc_task, signaling_task], return_when=asyncio.FIRST_COMPLETED)
295
296         for task in done:
297             task.result()
298         for task in pending:
299             task.cancel()
300     except OSError as e:
301         print(e)
302
303
304 def main():
305     logging.basicConfig(level=logging.DEBUG, format='%(asctime)-15s %(message)s')
306     default_source = 'wss://localhost:1234/ws_connect?id=cug'
307     default_dest = 'rtmp://localhost:1935/cug'
308     parser = argparse.ArgumentParser()
309     parser.add_argument('-s', '--source', default=default_source,
310                         help=f'Laplace signalling websocket URI, default: {default_source}')
311     parser.add_argument('-d', '--destination', default=default_dest,
312                         help=f'RTMP server URI, default: {default_dest}')
313     args = parser.parse_args()
314
315     Gst.init(None)
316     asyncio.run(run(args.source, args.destination))
317
318
319 if __name__ == '__main__':
320     main()