feat: add authorization token

app/crud.py

@@ -229,11 +229,13 @@ def _check_item_in_election(
         )
 
 
-def update_ballot(db: Session, ballot: schemas.BallotUpdate) -> schemas.BallotGet:
+def update_ballot(
+    db: Session, ballot: schemas.BallotUpdate, token: str
+) -> schemas.BallotGet:
     if ballot.votes == []:
         raise errors.BadRequestError("The ballot contains no vote")
 
-    payload = jws_verify(ballot.token)
+    payload = jws_verify(token)
     election_ref = payload["election"]
     vote_ids: list[int] = list(set(payload["votes"]))
 

app/main.py

@@ -1,6 +1,6 @@
 import typing as t
 import json
-from fastapi import Depends, FastAPI, HTTPException, Request, Body
+from fastapi import Depends, FastAPI, HTTPException, Request, Body, Header
 from fastapi.responses import JSONResponse, PlainTextResponse
 from fastapi.middleware.cors import CORSMiddleware
 from sqlalchemy.orm import Session
@@ -85,22 +85,22 @@ def create_ballot(
     ballot: schemas.BallotCreate,
     db: Session = Depends(get_db),
 ):
-    try:
-        return crud.create_ballot(db=db, ballot=ballot)
-    except JWSError:
-        raise errors.UnauthorizedError("Unverified token")
+    return crud.create_ballot(db=db, ballot=ballot)
 
 
 @app.put("/ballots", response_model=schemas.BallotGet)
-def update_ballot(vote: schemas.BallotUpdate, db: Session = Depends(get_db)):
-    try:
-        return crud.update_ballot(db=db, ballot=vote)
-    except JWSError:
-        raise errors.UnauthorizedError("Unverified token")
+def update_ballot(
+    ballot: schemas.BallotUpdate,
+    authorization: str = Header(),
+    db: Session = Depends(get_db),
+):
+    token = authorization.split("Bearer ")[1]
+    return crud.update_ballot(db=db, ballot=ballot, token=token)
 
 
-@app.get("/ballots/{token}", response_model=schemas.BallotGet)
-def get_ballot(token: str, db: Session = Depends(get_db)):
+@app.get("/ballots/", response_model=schemas.BallotGet)
+def get_ballot(authorization: str = Header(), db: Session = Depends(get_db)):
+    token = authorization.split("Bearer ")[1]
     return crud.get_ballot(db=db, token=token)
 
 

app/schemas.py

@@ -105,11 +105,6 @@ class VoteCreate(BaseModel):
         orm_mode = True
 
 
-class BallotUpdate(BaseModel):
-    token: str
-    votes: list[VoteCreate]
-
-
 def _in_a_long_time() -> datetime:
     """
     Provides the date in the future
@@ -222,3 +217,7 @@ class BallotGet(BaseModel):
 class BallotCreate(BaseModel):
     votes: list[VoteCreate]
     election_ref: str
+
+
+class BallotUpdate(BaseModel):
+    votes: list[VoteCreate]

app/tests/test_api.py

@@ -159,10 +159,12 @@ def test_create_ballot():
     token = data["token"]
 
     # Now, we check that we need the righ token to read the votes
-    response = client.get(f"/ballots/{token}WRONG")
+    response = client.get(
+        f"/ballots/", headers={"Authorization": f"Bearer {token}WRONG"}
+    )
     assert response.status_code == 401, response.text
 
-    response = client.get(f"/ballots/{token}")
+    response = client.get(f"/ballots/", headers={"Authorization": f"Bearer {token}"})
     assert response.status_code == 200, response.text
     data = response.json()
     for v1, v2 in zip(votes, data["votes"]):
@@ -188,7 +190,8 @@ def test_cannot_create_vote_on_restricted_election():
     # We create votes using the ID
     votes = _generate_votes_from_response("id", data)
     response = client.post(
-        f"/ballots", json={"votes": votes, "election_ref": election_ref}
+        f"/ballots",
+        json={"votes": votes, "election_ref": election_ref},
     )
     data = response.json()
     assert response.status_code == 400, data
@@ -220,7 +223,11 @@ def test_can_vote_on_restricted_election():
         {"candidate_id": candidate["id"], "grade_id": grade_id}
         for candidate in data["candidates"]
     ]
-    response = client.put(f"/ballots", json={"votes": votes, "token": token})
+    response = client.put(
+        f"/ballots",
+        json={"votes": votes},
+        headers={"Authorization": f"Bearer {token}"},
+    )
     data = response.json()
     assert response.status_code == 200, data