138 lines
3.2 KiB
Python
138 lines
3.2 KiB
Python
|
#!/usr/bin/env -S python3 -u
|
||
|
|
# Copyright 2019 The Matrix.org Foundation CIC
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import argparse
|
||
|
|
import socket
|
||
|
|
import json
|
||
|
|
from random import choice
|
||
|
|
from aiohttp import web
|
||
|
|
|
||
|
|
# The browsers ban some known ports, the dynamic port range doesn't contain any
|
||
|
|
# banned ports, so we use that.
|
||
|
|
port_range = range(49152, 65535)
|
||
|
|
|
||
|
|
shutdown_task = None
|
||
|
|
|
||
|
|
|
||
|
|
def to_weechat(message):
|
||
|
|
print(json.dumps(message))
|
||
|
|
|
||
|
|
|
||
|
|
async def get_token(request):
|
||
|
|
global shutdown_task
|
||
|
|
|
||
|
|
async def shutdown():
|
||
|
|
await asyncio.sleep(1)
|
||
|
|
raise KeyboardInterrupt
|
||
|
|
|
||
|
|
token = request.query.get("loginToken")
|
||
|
|
|
||
|
|
if not token:
|
||
|
|
raise KeyboardInterrupt
|
||
|
|
|
||
|
|
message = {
|
||
|
|
"type": "token",
|
||
|
|
"loginToken": token
|
||
|
|
}
|
||
|
|
|
||
|
|
# Send the token to weechat.
|
||
|
|
to_weechat(message)
|
||
|
|
# Initiate a shutdown.
|
||
|
|
shutdown_task = asyncio.ensure_future(shutdown())
|
||
|
|
# Respond to the browser.
|
||
|
|
return web.Response(text="Continuing in Weechat.")
|
||
|
|
|
||
|
|
|
||
|
|
def bind_socket(port=None):
|
||
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
|
|
|
||
|
|
if port is not None and port != 0:
|
||
|
|
sock.bind(("localhost", port))
|
||
|
|
return sock
|
||
|
|
|
||
|
|
while True:
|
||
|
|
port = choice(port_range)
|
||
|
|
|
||
|
|
try:
|
||
|
|
sock.bind(("localhost", port))
|
||
|
|
except OSError:
|
||
|
|
continue
|
||
|
|
|
||
|
|
return sock
|
||
|
|
|
||
|
|
|
||
|
|
async def wait_for_shutdown_task(_):
|
||
|
|
if not shutdown_task:
|
||
|
|
return
|
||
|
|
|
||
|
|
try:
|
||
|
|
await shutdown_task
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description="Start a web server that waits for a SSO token to be "
|
||
|
|
"passed with a GET request"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"-p", "--port",
|
||
|
|
help=("the port that the web server will be listening on, if 0 a "
|
||
|
|
"random port should be chosen"
|
||
|
|
),
|
||
|
|
type=int,
|
||
|
|
default=0
|
||
|
|
)
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
app = web.Application()
|
||
|
|
app.add_routes([web.get('/', get_token)])
|
||
|
|
|
||
|
|
if not 0 <= args.port <= 65535:
|
||
|
|
raise ValueError("Port needs to be 0-65535")
|
||
|
|
|
||
|
|
try:
|
||
|
|
sock = bind_socket(args.port)
|
||
|
|
except OSError as e:
|
||
|
|
message = {
|
||
|
|
"type": "error",
|
||
|
|
"message": str(e),
|
||
|
|
"code": e.errno
|
||
|
|
}
|
||
|
|
to_weechat(message)
|
||
|
|
return
|
||
|
|
|
||
|
|
host, port = sock.getsockname()
|
||
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||
|
|
|
||
|
|
message = {
|
||
|
|
"type": "redirectUrl",
|
||
|
|
"host": host,
|
||
|
|
"port": port
|
||
|
|
}
|
||
|
|
|
||
|
|
to_weechat(message)
|
||
|
|
|
||
|
|
app.on_shutdown.append(wait_for_shutdown_task)
|
||
|
|
web.run_app(app, sock=sock, handle_signals=True, print=None)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|