Skip to content

Commit

Permalink
Remove python-jose and add test cases
Browse files Browse the repository at this point in the history
Signed-off-by: Aisuko <[email protected]>
  • Loading branch information
Aisuko committed Jul 14, 2024
1 parent bf1e777 commit 2263eb9
Show file tree
Hide file tree
Showing 8 changed files with 142 additions and 52 deletions.
2 changes: 1 addition & 1 deletion .devcontainer/devcontainer.json
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@
"initializeCommand": "make prepare",

// Uncomment the next line to run commands after the container is created.
"postCreateCommand": "make install"
"postCreateCommand": "make install-dev"

// Configure tool-specific properties.
// "customizations": {},
Expand Down
13 changes: 8 additions & 5 deletions .vscode/settings.json
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
{
"python.testing.pytestArgs": [
"backend"
"python.testing.unittestArgs": [
"-v",
"-s",
"./backend",
"-p",
"test_*.py"
],
"python.testing.unittestEnabled": false,
"python.testing.pytestEnabled": true,
"DockerRun.DisableDockerrc": true
"python.testing.pytestEnabled": false,
"python.testing.unittestEnabled": true
}
12 changes: 12 additions & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -228,3 +228,15 @@ plugin:
.PHONY: expo
expo:
@poetry -C backend export -f requirements.txt --output backend/requirements.txt

############################################################################################################
# Linter and test

.PHONY: lint
lint:
@ruff check --output-format=github .


.PHONY: test
test:
@poetry run -C backend python -m unittest discover -s backend/tests/ -v
5 changes: 3 additions & 2 deletions backend/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,15 @@ pydantic-settings = "2.2.1"
pymilvus = "2.3.7"
python-decouple = "3.8"
python-dotenv = "1.0.1"
python-jose = "3.3.0"
PyJWT= "2.8.0"
python-multipart = "0.0.9"
python-slugify = "8.0.4"
sqlalchemy = "2.0.29"
trio = "0.25.0"
uvicorn = "0.29.0"
openai = "1.35.7"


[tool.poetry.dev-dependencies]
python = "^3.11"
alembic = "1.13.1"
Expand All @@ -56,7 +57,7 @@ pydantic-settings = "2.2.1"
pymilvus = "2.3.7"
python-decouple = "3.8"
python-dotenv = "1.0.1"
python-jose = "3.3.0"
PyJWT= "2.8.0"
python-multipart = "0.0.9"
python-slugify = "8.0.4"
sqlalchemy = "2.0.29"
Expand Down
21 changes: 6 additions & 15 deletions backend/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -308,9 +308,6 @@ distro==1.9.0 ; python_version >= "3.11" and python_version < "4.0" \
dnspython==2.6.1 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:5ef3b9680161f6fa89daf8ad451b5f1a33b18ae8a1c6778cdf4b43f08c0a6e50 \
--hash=sha256:e8f0f9c23a7b7cb99ded64e6c3a6f3e701d78f50c55e002b839dea7225cff7cc
ecdsa==0.19.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:2cea9b88407fdac7bbeca0833b189e4c9c53f2ef1e1eaa29f6224dbc809b707a \
--hash=sha256:60eaad1199659900dd0af521ed462b793bbdf867432b3948e87416ae4caf6bf8
email-validator==2.1.1 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:200a70680ba08904be6d1eef729205cc0d687634399a5924d842533efb824b84 \
--hash=sha256:97d882d174e2a65732fb43bfce81a3a834cbc1bde8bf419e30ef5ea976370a05
Expand Down Expand Up @@ -735,9 +732,6 @@ pyarrow==16.1.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:f68f409e7b283c085f2da014f9ef81e885d90dcd733bd648cfba3ef265961848 \
--hash=sha256:fbef391b63f708e103df99fbaa3acf9f671d77a183a07546ba2f2c297b361e83 \
--hash=sha256:febde33305f1498f6df85e8020bca496d0e9ebf2093bab9e0f65e2b4ae2b3444
pyasn1==0.6.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:3a35ab2c4b5ef98e17dfdec8ab074046fbda76e281c5a706ccd82328cfc8f64c \
--hash=sha256:cca4bb0f2df5504f02f6f8a775b6e416ff9b0b3b16f7ee80b5a3153d9b804473
pycparser==2.22 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6 \
--hash=sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc
Expand Down Expand Up @@ -870,6 +864,9 @@ pydantic-settings==2.2.1 ; python_version >= "3.11" and python_version < "4.0" \
pydantic==2.8.2 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:6f62c13d067b0755ad1c21a34bdd06c0c12625a22b0fc09c6b149816604f7c2a \
--hash=sha256:73ee9fddd406dc318b885c7a2eab8a6472b68b8fb5ba8150949fc3db939f23c8
pyjwt==2.8.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:57e28d156e3d5c10088e0c68abb90bfac3df82b40a71bd0daa20c65ccd5c23de \
--hash=sha256:59127c392cc44c2da5bb3192169a91f429924e17aff6534d70fdc02ab3e04320
pymilvus==2.3.7 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:37d5a360d671c6fe23fe1dd4e6b41af6e4b6d6488ad8e43a06afe23d02f98272 \
--hash=sha256:b8df5b8db3a82209c33b7211e0b9ef4a63ee00cb2976ccb1e9f5b92a2c2d5b82
Expand All @@ -882,9 +879,6 @@ python-decouple==3.8 ; python_version >= "3.11" and python_version < "4.0" \
python-dotenv==1.0.1 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:e324ee90a023d808f1959c46bcbc04446a10ced277783dc6ee09987c37ec10ca \
--hash=sha256:f7b63ef50f1b690dddf550d03497b66d609393b40b564ed0d674909a68ebf16a
python-jose==3.3.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:55779b5e6ad599c6336191246e95eb2293a9ddebd555f796a65f838f07e5d78a \
--hash=sha256:9b1376b023f8b298536eedd47ae1089bcdb848f1535ab30555cd92002d78923a
python-multipart==0.0.9 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:03f54688c663f1b7977105f021043b0793151e4cb1c1a9d4a11fc13d622c4026 \
--hash=sha256:97ca7b8ea7b05f977dc3849c3ba99d51689822fab725c3703af7c866a0c2b215
Expand All @@ -897,12 +891,9 @@ pytz==2024.1 ; python_version >= "3.11" and python_version < "4.0" \
requests==2.32.3 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \
--hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
rsa==4.9 ; python_version >= "3.11" and python_version < "4" \
--hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \
--hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21
setuptools==70.2.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:b8b8060bb426838fbe942479c90296ce976249451118ef566a5a0b7d8b78fb05 \
--hash=sha256:bd63e505105011b25c3c11f753f7e3b8465ea739efddaccef8f0efac2137bac1
setuptools==70.3.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:f171bab1dfbc86b132997f26a119f6056a57950d058587841a0082e8830f9dc5 \
--hash=sha256:fe384da74336c398e0d956d1cae0669bc02eed936cdb1d49b57de1990dc11ffc
six==1.16.0 ; python_version >= "3.11" and python_version < "4.0" \
--hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
--hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
Expand Down
34 changes: 23 additions & 11 deletions backend/src/securities/authorizations/jwt.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,28 @@
# coding=utf-8

# Copyright [2024] [SkywardAI]
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime

import pydantic
from jose import jwt as jose_jwt, JWTError as JoseJWTError
import jwt as pyjwt
from fastapi import Request, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

from src.config.manager import settings
from src.models.db.account import Account
from src.models.schemas.jwt import JWTAccount, JWToken
from src.utilities.exceptions.database import EntityDoesNotExist

class JWTGenerator:
def __init__(self):
Expand All @@ -23,30 +37,28 @@ def _generate_jwt_token(
to_encode = jwt_data.copy()

if expires_delta:
expire = datetime.datetime.utcnow() + expires_delta
expire = datetime.datetime.now(datetime.UTC) + expires_delta

else:
expire = datetime.datetime.utcnow() + datetime.timedelta(minutes=settings.JWT_MIN)
expire = datetime.datetime.now(datetime.UTC) + datetime.timedelta(minutes=settings.JWT_MIN)

to_encode.update(JWToken(exp=expire, sub=settings.JWT_SUBJECT).dict())
to_encode.update(JWToken(exp=expire, sub=settings.JWT_SUBJECT).model_dump())

return jose_jwt.encode(to_encode, key=settings.JWT_SECRET_KEY, algorithm=settings.JWT_ALGORITHM)
return pyjwt.encode(to_encode, key=settings.JWT_SECRET_KEY, algorithm=settings.JWT_ALGORITHM)

def generate_access_token(self, account: Account) -> str:
if not account:
raise EntityDoesNotExist("Cannot generate JWT token for without Account entity!")

return self._generate_jwt_token(
jwt_data=JWTAccount(username=account.username, email=account.email).dict(), # type: ignore
jwt_data=JWTAccount(username=account.username, email=account.email).model_dump(), # type: ignore
expires_delta=datetime.timedelta(minutes=settings.JWT_ACCESS_TOKEN_EXPIRATION_TIME),
)

def retrieve_details_from_token(self, token: str) -> dict:
try:
payload = jose_jwt.decode(token=token, key=settings.JWT_SECRET_KEY, algorithms=[settings.JWT_ALGORITHM])
payload = pyjwt.decode(token, key=settings.JWT_SECRET_KEY, algorithms=[settings.JWT_ALGORITHM])
jwt_account = JWTAccount(username=payload["username"], email=payload["email"])

except JoseJWTError as token_decode_error:
except pyjwt.exceptions.DecodeError as token_decode_error:
raise ValueError("Unable to decode JWT Token") from token_decode_error

except pydantic.ValidationError as validation_error:
Expand Down
58 changes: 58 additions & 0 deletions backend/tests/unit_tests/test_jwt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
# coding=utf-8

# Copyright [2024] [SkywardAI]
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import jwt as pyjwt

@unittest.skip("Skip this test, it is a evidence of a security vulnerability")
class TestJWTReplacedSolution(unittest.TestCase):

JWT_SECRET_KEY="YOUR-KEY"
ALGORITHM="HS256"
content={"some": "payload"}

@classmethod
def setUpClass(cls) -> None:
return super().setUpClass()

@classmethod
def tearDownClass(cls) -> None:
return super().tearDownClass()


def test_jwt_jose(self):
pass
# jose_str=jose_jwt.encode(self.content, key=self.JWT_SECRET_KEY, algorithm=self.ALGORITHM)

# pyjwt_str=pyjwt.encode(self.content, key=self.JWT_SECRET_KEY, algorithm=self.ALGORITHM)

# # jose and pyjwt should produce the same token
# assert jose_str == pyjwt_str

def test_jwt_decode(self):
pass
# jose_str=jose_jwt.encode(self.content, key=self.JWT_SECRET_KEY, algorithm=self.ALGORITHM)

# pyjwt_str=pyjwt.encode(self.content, key=self.JWT_SECRET_KEY, algorithm=self.ALGORITHM)

# # jose and pyjwt should produce the same token
# assert jose_str == pyjwt_str

# # decode the token
# jose_decoded=jose_jwt.decode(jose_str, key=self.JWT_SECRET_KEY, algorithms=[self.ALGORITHM])
# pyjwt_decoded=pyjwt.decode(pyjwt_str, key=self.JWT_SECRET_KEY, algorithms=[self.ALGORITHM])

# # jose and pyjwt should produce the same decoded token
# assert jose_decoded == pyjwt_decoded
49 changes: 31 additions & 18 deletions backend/tests/unit_tests/test_method_kit.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,28 +13,41 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from src.utillities httpkit import MethodKit
import unittest
from backend.src.utilities.httpkit.method_kit import MethodKit


def test_http_post()-> None:
"""
Test the http_post method.
In this test case we calculate the the length of tokens of the content "Hello, World!".
"""
@unittest.skip("Skip due to the fact that the server is not running")
class TestHTTPKits(unittest.TestCase):
url = "http://llamacpp:8080/tokenize"
jason_content = {"content": "Hello, World!"}
headers={'Content-Type': 'application/json'}
timeout = 10

res = MethodKit.http_post(
url=url,
jason_content=jason_content,
headers=headers,
timeout=timeout
)

assert res.status_code == 200
# length of token is a integer
assert isinstance(res.json().get('tokens'), int)
@classmethod
def setUpClass(cls) -> None:
return super().setUpClass()

@classmethod
def tearDownClass(cls) -> None:
return super().tearDownClass()


def test_http_post(self)-> None:
"""
Test the http_post method.
In this test case we calculate the the length of tokens of the content "Hello, World!".
"""

res = MethodKit.http_post(
url=self.url,
jason_content=self.jason_content,
headers=self.headers,
timeout=self.timeout
)

assert res.status_code == 200
# length of token is a integer
assert isinstance(res.json().get('tokens'), int)

0 comments on commit 2263eb9

Please sign in to comment.