diff --git a/AuthToken.py b/AuthToken.py new file mode 100644 index 0000000..e743c5a --- /dev/null +++ b/AuthToken.py @@ -0,0 +1,87 @@ +import os +import json +import click +import time +import requests + +class AuthToken: + def __init__(self): + self.tokenfile = "token.json" + + def save_oauth_token(self, token): + tokenfile_content = {} + # check if file extists + if os.path.exists(self.tokenfile): + with open(self.tokenfile, "r") as f: + tokenfile_content = json.load(f) + tokenfile_content["oauth_token"] = token + with open(self.tokenfile, "w") as f: + json.dump(tokenfile_content, f) + + def get_oauth_token(self): + # check if token in tokenfile + if os.path.exists(self.tokenfile): + with open(self.tokenfile, "r") as f: + tokendict = json.load(f) + if "oauth_token" in tokendict: + return tokendict["oauth_token"] + + # if not, check if the github copilot extension is installed on the system + path = os.path.expanduser("~/.config/github-copilot/hosts.json") + if os.path.exists(path): + with open(path, "r") as f: + hosts = json.load(f) + if "github.com" in hosts: + token = hosts["github.com"]["oauth_token"] + self.save_oauth_token(token) + return token + + # if not, ask the user to enter the token + click.echo("Please enter your github access token:") + token = click.prompt("Token", type=str) + self.save_oauth_token(token) + return token + + def save_github_api_token(self, token: str, expires_at): + tokenfile_content = {} + # check if file extists + if os.path.exists(self.tokenfile): + with open(self.tokenfile, "r") as f: + tokenfile_content = json.load(f) + tokenfile_content["token"] = token + tokenfile_content["expires_at"] = expires_at + with open(self.tokenfile, "w") as f: + json.dump(tokenfile_content, f) + + def request_github_api_token(self): + # get the oauth token + oauth_token = self.get_oauth_token() + + headers = { + 'Authorization': f"Bearer {oauth_token}", + } + response = requests.get('https://api.github.com/copilot_internal/v2/token', headers=headers) + + if response.status_code != 200: + return None + + token = response.json() + self.save_github_api_token(token["token"], token["expires_at"]) + return token["token"] + + + def get_github_api_token(self): + # check if token in tokenfile + if os.path.exists(self.tokenfile): + with open(self.tokenfile, "r") as f: + tokendict = json.load(f) + if "token" in tokendict and "expires_at" in tokendict: + if tokendict["expires_at"] > time.time(): + return tokendict["token"] + + # request new api token + return self.request_github_api_token() + + + + diff --git a/copilot.py b/copilot.py index 9c6a100..81060cb 100755 --- a/copilot.py +++ b/copilot.py @@ -4,6 +4,7 @@ import os import sys import time from os.path import expanduser +from AuthToken import AuthToken import click from rich import print from rich.panel import Panel @@ -14,50 +15,6 @@ try: except ImportError: import readline - - -# get a new github copilot token -def get_token(): - - valid = False - token = {} - - # check if a token exists in the file - if os.path.exists("token.json"): - with open("token.json", "r") as f: - token = json.load(f) - - # check if the token is valid - if "token" in token and "expires_at" in token: - if token["expires_at"] > time.time(): - valid = True - - if not valid: - # read token from file - - # read file from home directory - home = expanduser("~") - with open(f"{home}/.config/github-copilot/hosts.json", "r") as f: - terms = json.load(f) - auth_token = terms["github.com"]["oauth_token"] - - - headers = { - 'Authorization': f"Bearer {auth_token}", - } - response = requests.get('https://api.github.com/copilot_internal/v2/token', headers=headers) - - if response.status_code != 200: - return None - - token = response.json() - - # save the token in a json file with the expiration time - with open("token.json", "w") as f: - json.dump(token, f) - - return token["token"] - # get a suggestion using github copilot def get_suggestion(prompt, token): headers = { @@ -125,7 +82,9 @@ def main(): def generate_suggestion(prompt): prompt = " ".join(prompt) prompt = "!/bin/bash\n\n" + prompt + ":\n" - token = get_token() + + authToken = AuthToken() + token = authToken.get_github_api_token() suggestion = get_suggestion(prompt, token) panel = Panel(Text(suggestion, justify="center"), title_align="center", title="Execute? \[y/n/e/c]", expand=False) diff --git a/setup.py b/setup.py index 5cc6eb6..a226ab5 100644 --- a/setup.py +++ b/setup.py @@ -9,7 +9,7 @@ setup( author='Victor Mylle', author_email='victor.mylle@hotmail.com', license='MIT', - py_modules=['copilot'], + py_modules=['copilot', 'AuthToken'], install_requires=[ 'requests', 'click', diff --git a/token.json b/token.json deleted file mode 100644 index 1554859..0000000 --- a/token.json +++ /dev/null @@ -1 +0,0 @@ -{"expires_at": 1669988057, "public_suggestions": "disabled", "refresh_in": 1500, "sku": "free_educational", "telemetry": "disabled", "token": "tid=d47460691ef58366f713e508d59d807b;exp=1669988057;sku=free_educational:8f88da8df463706243fcf52df0f6e9d5ec3ca0126ccdc919c2528ce4f7573101", "tracking_id": "d47460691ef58366f713e508d59d807b"} \ No newline at end of file