Spaces:
Runtime error
Runtime error
File size: 4,033 Bytes
105b369 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Optional
from phi.cli.settings import phi_cli_settings
class CliAuthRequestHandler(BaseHTTPRequestHandler):
"""Request Handler to accept the CLI auth token after the web based auth flow.
References:
https://medium.com/@hasinthaindrajee/browser-sso-for-cli-applications-b0be743fa656
https://gist.github.com/mdonkers/63e115cc0c79b4f6b8b3a6b797e485c7
TODO:
* Fix the header and limit to only localhost or phidata.com
"""
def _set_response(self):
self.send_response(200)
self.send_header("Content-type", "application/json")
self.send_header("Access-Control-Allow-Origin", "*")
self.send_header("Access-Control-Allow-Headers", "*")
self.send_header("Access-Control-Allow-Methods", "POST")
self.end_headers()
# def do_GET(self):
# logger.info("GET request,\nPath: %s\nHeaders:\n%s\n", str(self.path), str(self.headers))
# self._set_response()
# self.wfile.write("GET request for {}".format(self.path).encode('utf-8'))
def do_OPTIONS(self):
# logger.debug(
# "OPTIONS request,\nPath: %s\nHeaders:\n%s\n",
# str(self.path),
# str(self.headers),
# )
self._set_response()
# self.wfile.write("OPTIONS request for {}".format(self.path).encode('utf-8'))
def do_POST(self):
content_length = int(self.headers["Content-Length"]) # <--- Gets the size of data
post_data = self.rfile.read(content_length) # <--- Gets the data itself
decoded_post_data = post_data.decode("utf-8")
# logger.debug(
# "POST request,\nPath: {}\nHeaders:\n{}\n\nBody:\n{}\n".format(
# str(self.path), str(self.headers), decoded_post_data
# )
# )
# logger.debug("Data: {}".format(decoded_post_data))
# logger.info("type: {}".format(type(post_data)))
phi_cli_settings.tmp_token_path.touch(exist_ok=True)
phi_cli_settings.tmp_token_path.write_text(decoded_post_data)
# TODO: Add checks before shutting down the server
self.server.running = False # type: ignore
self._set_response()
def log_message(self, format, *args):
pass
class CliAuthServer:
"""
Source: https://stackoverflow.com/a/38196725/10953921
"""
def __init__(self, port: int = 9191):
import threading
self._server = HTTPServer(("", port), CliAuthRequestHandler)
self._thread = threading.Thread(target=self.run)
self._thread.daemon = True
self._server.running = False # type: ignore
def run(self):
self._server.running = True # type: ignore
while self._server.running: # type: ignore
self._server.handle_request()
def start(self):
self._thread.start()
def shut_down(self):
self._thread.close() # type: ignore
def check_port(port: int):
import socket
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
try:
return s.connect_ex(("localhost", port)) == 0
except Exception as e:
print(f"Error occurred: {e}")
return False
def get_port_for_auth_server():
starting_port = 9191
for port in range(starting_port, starting_port + 100):
if not check_port(port):
return port
def get_auth_token_from_web_flow(port: int) -> Optional[str]:
"""
GET request: curl http://localhost:9191
POST request: curl -d "foo=bar&bin=baz" http://localhost:9191
"""
import json
server = CliAuthServer(port)
server.run()
if phi_cli_settings.tmp_token_path.exists() and phi_cli_settings.tmp_token_path.is_file():
auth_token_str = phi_cli_settings.tmp_token_path.read_text()
auth_token_json = json.loads(auth_token_str)
phi_cli_settings.tmp_token_path.unlink()
return auth_token_json.get("AuthToken", None)
return None
|