diff --git a/pyros_api.py b/pyros_api.py index f0ff158..89ad5ac 100755 --- a/pyros_api.py +++ b/pyros_api.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 -from email import header import os +from urllib3.exceptions import ProtocolError +from requests.exceptions import ConnectionError import requests import sys import yaml @@ -12,10 +13,9 @@ class PyrosAPI: """ Request Pyros API with an PyrosUser account linked to the username and password """ - BASE_PYROS_URL = "http://localhost:8000/api/" - #BASE_PYROS_URL = "http://pyros.irap.omp.eu/api/" + #host = "http://pyros.irap.omp.eu/api/" - def __init__(self) -> None: + def __init__(self, host) -> None: """ Initialize PyrosAPI class by getting the authentification token required to use the API @@ -23,6 +23,11 @@ class PyrosAPI: username : Username of the user (mail adress) password : Password of the user """ + if host == None: + self.host = "http://localhost:8000" + else: + self.host = host + self.host = self.host+"/api/" self.get_token() def get_token(self): @@ -38,23 +43,27 @@ class PyrosAPI: self.token = token_file.read() return else: - url = f"{self.BASE_PYROS_URL}api-token-auth/" + url = f"{self.host}api-token-auth/" print("Token needs to be generated, you need to authentificate to PyROS :") username = input("Enter your username: ") - try: password = getpass.getpass() except Exception as error: print('ERROR', error) - request = requests.post( - url, data={"username": username, "password": password}) - if request.status_code != 200: - print("Authentification failed, please retry") - exit(1) - json_response = request.json() - self.token = json_response["token"] - with open("TOKEN", "w+") as token_file: - token_file.write(self.token) + try: + request = requests.post( + url, data={"username": username, "password": password}) + except (ConnectionError, ConnectionResetError, ProtocolError): + print( + f"Server {self.host} doesn't respond. The website might be down.") + exit(1) + if request.status_code != 200: + print("Authentification failed, please retry") + exit(1) + json_response = request.json() + self.token = json_response["token"] + with open("TOKEN", "w+") as token_file: + token_file.write(self.token) def get_url(self, url: str) -> dict: """ @@ -67,8 +76,12 @@ class PyrosAPI: dict : Json object that represents the response of the API """ headers = {"Authorization": f"Token {self.token}"} - - request = requests.get(self.BASE_PYROS_URL+url, headers=headers) + try: + request = requests.get(self.host+url, headers=headers) + except (ConnectionError, ConnectionResetError, ProtocolError): + print( + f"Server {self.host} doesn't respond. The website might be down.") + exit(1) return request.json() def submit_sequence_file(self, file: str) -> dict: @@ -85,12 +98,17 @@ class PyrosAPI: "Authorization": f"Token {self.token}", 'Content-type': 'application/json', } - url = f"{self.BASE_PYROS_URL}submit_sequence" + url = f"{self.host}submit_sequence" yaml_file = open(file, "r") sequence_file_as_dict = yaml.safe_load(yaml_file) print(f"File content : {sequence_file_as_dict}") - request = requests.put(url, headers=headers, - json=sequence_file_as_dict) + try: + request = requests.put(url, headers=headers, + json=sequence_file_as_dict) + except (ConnectionError, ConnectionResetError, ProtocolError): + print( + f"Server {self.host} doesn't respond. The website might be down.") + exit(1) print(f"Request status code {request.status_code}") return request.json() @@ -103,23 +121,36 @@ class PyrosAPI: Returns: dict : Json object that represents the response of the API """ - url = f"{self.BASE_PYROS_URL}full_sequences/get_sequences_for_period/" + url = f"{self.host}full_sequences/get_sequences_for_period/" headers = {"Authorization": f"Token {self.token}"} - response = requests.get( - url, params={"start_date": start_date, "end_date": end_date}, headers=headers) + try: + response = requests.get( + url, params={"start_date": start_date, "end_date": end_date}, headers=headers) + except (ConnectionError, ConnectionResetError, ProtocolError): + print( + f"Server {self.host} doesn't respond. The website might be down.") + exit(1) return response.json() def logout(self): - url = f"{self.BASE_PYROS_URL}logout/" + url = f"{self.host}logout/" header = {"Authorization": f"Token {self.token}"} - response = requests.get(url, headers=header) + try: + response = requests.get(url, headers=header) + except (ConnectionError, ConnectionResetError, ProtocolError): + print( + f"Server {self.host} doesn't respond. The website might be down.") + exit(1) return response.content.decode("utf-8") @click.group() +@click.option('--host', '-h', help='host name (example: http://localhost:8000 or http://pyros.omp.eu) without last slash') @click.pass_context -def cli(ctx): - ctx.obj = PyrosAPI() +def cli(ctx, host): + if not host: + host = None + ctx.obj = PyrosAPI(host) @cli.command("get_sequences_for_period", help="Get sequences for a period of date (start/end)") -- libgit2 0.21.2