Skip to content

Commit 8c55bab

Browse files
committed
Require user to enter own GitHub token
1 parent f329106 commit 8c55bab

File tree

19 files changed

+242
-23
lines changed

19 files changed

+242
-23
lines changed

.env.example

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
export OPEN_AI_KEY=fill_me_in
2-
export GITHUB_TOKEN=fill_me_in
32
export USE_FLASK_DEBUG_MODE=true
3+
export FLASK_SECRET_KEY=e0bc47d4a5c3eedd2a5aa2dc0d6c1b0454f523241467f198853e4057a418d1f3

.github/workflows/pipeline.yml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,5 +54,4 @@ jobs:
5454
region: us-central1
5555
flags: --allow-unauthenticated --platform managed
5656
secrets: |
57-
GITHUB_TOKEN=GITHUB_TOKEN:latest
5857
OPEN_AI_KEY=OPEN_AI_KEY:latest

discovery/app.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,21 +3,22 @@
33
from flask import Flask
44
from openai import OpenAI
55

6+
from discovery.authentication_page import authentication_page
67
from discovery.environment import Environment
7-
from discovery.github_support.github_client import GithubClient
88
from discovery.index_page import index_page
9-
from discovery.repository_agent.repository_agent import create_repository_agent
9+
from discovery.repository_agent.repository_agent import repository_agent_creator
1010

1111
logger = logging.getLogger(__name__)
1212

1313

1414
def create_app(env: Environment = Environment.from_env()):
1515
app = Flask(__name__)
16-
open_ai_client = OpenAI(api_key=env.open_ai_key)
17-
github_client = GithubClient(access_token=env.github_token)
16+
app.secret_key = env.flask_secret_key
1817

19-
agent = create_repository_agent(open_ai_client, github_client)
18+
open_ai_client = OpenAI(api_key=env.open_ai_key)
19+
agent_creator = repository_agent_creator(open_ai_client)
2020

21-
app.register_blueprint(index_page(agent))
21+
app.register_blueprint(index_page(agent_creator))
22+
app.register_blueprint(authentication_page())
2223

2324
return app

discovery/auth/__init__.py

Whitespace-only changes.
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
from functools import wraps
2+
3+
from flask import redirect, g
4+
5+
from discovery.auth.session_manager import SessionManager
6+
7+
8+
def require_authentication(func):
9+
@wraps(func)
10+
def wrapper(*args, **kwargs):
11+
user = SessionManager.user()
12+
if user is None:
13+
return redirect("/login")
14+
15+
g.username = user.username
16+
g.github_token = user.github_token
17+
return func(*args, **kwargs)
18+
19+
return wrapper

discovery/auth/session_manager.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
from dataclasses import dataclass
2+
from typing import Optional
3+
4+
from flask import session
5+
6+
7+
@dataclass
8+
class SessionUser:
9+
username: str
10+
github_token: str
11+
12+
13+
class SessionManager:
14+
@staticmethod
15+
def login(username: str, github_token: str):
16+
session["username"] = username
17+
session["github_token"] = github_token
18+
19+
@staticmethod
20+
def user() -> Optional[SessionUser]:
21+
if "username" not in session or "github_token" not in session:
22+
return None
23+
24+
return SessionUser(
25+
username=session["username"],
26+
github_token=session["github_token"],
27+
)
28+
29+
@staticmethod
30+
def logout():
31+
session.clear()

discovery/authentication_page.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import logging
2+
3+
from flask import Blueprint, render_template, request, Response, redirect, g
4+
from flask.typing import ResponseReturnValue
5+
6+
from discovery.auth.requre_authentication import require_authentication
7+
from discovery.auth.session_manager import SessionManager
8+
from discovery.github_support.github_client import GithubClient
9+
10+
logger = logging.getLogger(__name__)
11+
12+
13+
def authentication_page() -> Blueprint:
14+
page = Blueprint('authentication_page', __name__)
15+
16+
@page.get('/login')
17+
def login_page() -> ResponseReturnValue:
18+
return render_template('login.html')
19+
20+
@page.get('/logout')
21+
@require_authentication
22+
def log_out() -> ResponseReturnValue:
23+
logger.info(f"Logging out user %s", g.username)
24+
SessionManager.logout()
25+
return redirect("/login")
26+
27+
@page.post('/login')
28+
def login_with_token() -> ResponseReturnValue:
29+
github_token = request.form.get('token')
30+
github_client = GithubClient(github_token)
31+
user = github_client.get_user()
32+
if user is None:
33+
logger.error(f"User not found")
34+
return Response("User not found", status=400)
35+
36+
SessionManager.login(user.name, github_token)
37+
logger.info(f"Logged in user %s", user.name)
38+
39+
return redirect("/")
40+
41+
return page

discovery/environment.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,21 @@
55
@dataclass
66
class Environment:
77
open_ai_key: str
8-
github_token: str
98
root_log_level: str
109
discovery_log_level: str
1110
use_flask_debug_mode: bool
1211
port: str
12+
flask_secret_key: str
1313

1414
@classmethod
1515
def from_env(cls) -> 'Environment':
1616
return cls(
1717
open_ai_key=require_env('OPEN_AI_KEY'),
18-
github_token=require_env('GITHUB_TOKEN'),
1918
root_log_level=os.environ.get('ROOT_LOG_LEVEL', 'INFO'),
2019
discovery_log_level=os.environ.get('DISCOVERY_LOG_LEVEL', 'INFO'),
2120
use_flask_debug_mode=os.environ.get('USE_FLASK_DEBUG_MODE', 'false') == 'true',
2221
port=os.environ.get('PORT', '5050'),
22+
flask_secret_key=require_env('FLASK_SECRET_KEY'),
2323
)
2424

2525
def require_env(name: str) -> str:

discovery/github_support/github_client.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import logging
22
from dataclasses import dataclass
3-
from typing import List
3+
from typing import List, Optional
44

55
import requests
66

@@ -17,6 +17,11 @@ class Repository:
1717
forks: int
1818

1919

20+
@dataclass
21+
class GithubUser:
22+
name: str
23+
24+
2025
logger = logging.getLogger(__name__)
2126

2227

@@ -28,6 +33,14 @@ def __init__(self, access_token: str):
2833
"X-GitHub-Api-Version": "2022-11-28",
2934
}
3035

36+
def get_user(self) -> Optional[GithubUser]:
37+
response = requests.get(f"https://api.github.com/user", headers=self.request_headers)
38+
if response.status_code != 200:
39+
logger.error(f"Failed to get user")
40+
return None
41+
42+
return GithubUser(name=response.json()["login"])
43+
3144
def list_repositories_for_organization(self, organization: str) -> List[Repository]:
3245
response = requests.get(f"https://api.github.com/orgs/{organization}/repos", headers=self.request_headers)
3346
if response.status_code != 200:
@@ -37,7 +50,7 @@ def list_repositories_for_organization(self, organization: str) -> List[Reposito
3750
return [self.__repo_from_json(repo) for repo in response.json()]
3851

3952
def list_repositories_for_user(self, user: str) -> List[Repository]:
40-
response = requests.get(f"https://api.github.com/user/{user}/repos", headers=self.request_headers)
53+
response = requests.get(f"https://api.github.com/users/{user}/repos", headers=self.request_headers)
4154
if response.status_code != 200:
4255
logger.error(f"Failed to get repositories for user {user}: {response.text}")
4356
return []

discovery/index_page.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,32 @@
1-
from typing import Protocol
1+
from typing import Protocol, Callable
22

33
import mistune
4-
from flask import Blueprint, render_template, request
4+
from flask import Blueprint, render_template, request, g
55
from flask.typing import ResponseReturnValue
66

77
from discovery.agent_support.agent import AgentResult
8+
from discovery.auth.requre_authentication import require_authentication
9+
from discovery.github_support.github_client import GithubClient
810

911

1012
class AiAgent(Protocol):
1113
def answer(self, query: str) -> AgentResult:
1214
...
1315

1416

15-
def index_page(agent: AiAgent) -> Blueprint:
17+
def index_page(agent_creator: Callable[[GithubClient], AiAgent]) -> Blueprint:
1618
page = Blueprint('index_page', __name__)
1719

1820
@page.get('/')
21+
@require_authentication
1922
def index() -> ResponseReturnValue:
2023
return render_template('index.html')
2124

2225
@page.post('/')
26+
@require_authentication
2327
def query() -> ResponseReturnValue:
2428
user_query = request.form.get('query')
29+
agent = agent_creator(GithubClient(g.github_token))
2530
result = agent.answer(user_query)
2631

2732
return render_template(

0 commit comments

Comments
 (0)