From 808a0fa6fc3b96b17f59cab86fdcd60f7e12ca06 Mon Sep 17 00:00:00 2001 From: Saurab-Shrestha Date: Sun, 25 Feb 2024 09:54:56 +0545 Subject: [PATCH] Updated with audit table --- .../b7b23e5d5214_added_password_expiry.py | 32 +++++++++++++++++++ private_gpt/server/ingest/ingest_router.py | 6 ++-- private_gpt/users/api/deps.py | 8 +++++ private_gpt/users/api/v1/routers/auth.py | 4 +-- .../users/api/v1/routers/departments.py | 27 ++++++++++++++-- private_gpt/users/core/security.py | 2 +- private_gpt/users/crud/department_crud.py | 14 +++++++- private_gpt/users/models/user.py | 29 +++++++++++++++++ private_gpt/users/schemas/__init__.py | 2 +- private_gpt/users/schemas/department.py | 8 +++-- private_gpt/users/schemas/user.py | 6 +++- 11 files changed, 124 insertions(+), 14 deletions(-) create mode 100644 alembic/versions/b7b23e5d5214_added_password_expiry.py diff --git a/alembic/versions/b7b23e5d5214_added_password_expiry.py b/alembic/versions/b7b23e5d5214_added_password_expiry.py new file mode 100644 index 00000000..1100122f --- /dev/null +++ b/alembic/versions/b7b23e5d5214_added_password_expiry.py @@ -0,0 +1,32 @@ +"""added password expiry + +Revision ID: b7b23e5d5214 +Revises: 62c1a29320fc +Create Date: 2024-02-24 19:47:49.069873 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'b7b23e5d5214' +down_revision: Union[str, None] = '62c1a29320fc' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + # op.create_unique_constraint('unique_user_role', 'user_roles', ['user_id', 'role_id', 'company_id']) + op.add_column('users', sa.Column('password_created', sa.DateTime(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column('users', 'password_created') + # op.drop_constraint('unique_user_role', 'user_roles', type_='unique') + # ### end Alembic commands ### diff --git a/private_gpt/server/ingest/ingest_router.py b/private_gpt/server/ingest/ingest_router.py index 00b4fc2f..35bd3e98 100644 --- a/private_gpt/server/ingest/ingest_router.py +++ b/private_gpt/server/ingest/ingest_router.py @@ -220,7 +220,6 @@ def ingest_file( object="list", model="private-gpt", data=ingested_documents) log_audit(model='Document', action='create', details={ - 'status': '200', 'filename': file.filename, 'user': current_user.fullname, }, user_id=current_user.id) @@ -242,7 +241,6 @@ def ingest_file( async def common_ingest_logic( request: Request, - db: Session, ocr_file, current_user, @@ -274,10 +272,10 @@ async def common_ingest_logic( with open(upload_path, "wb") as f: f.write(file.read()) - file.seek(0) # Move the file pointer back to the beginning + file.seek(0) ingested_documents = service.ingest_bin_data(file_name, file) log_audit(model='Document', action='create', - details={'status': 200, 'message': "file uploaded successfully."}, user_id=current_user.id) + details={'status': "SUCCESS", 'message': f"{file_name} uploaded successfully."}, user_id=current_user.id) logger.info( f"{file_name} is uploaded by the {current_user.fullname}.") diff --git a/private_gpt/users/api/deps.py b/private_gpt/users/api/deps.py index 274b51a0..c42525c6 100644 --- a/private_gpt/users/api/deps.py +++ b/private_gpt/users/api/deps.py @@ -139,3 +139,11 @@ def get_audit_logger(request: Request, db: Session = Depends(get_db)): return lambda model, action, details, user_id=None: log_audit_entry(db, model, action, details, user_id) except Exception as e: raise HTTPException(status_code=500, detail=f"Error in get_audit_logger: {str(e)}") + + +def get_current_active_user( + current_user: models.User = Security(get_current_user, scopes=[],), +) -> models.User: + if not crud.user.is_active(current_user): + raise HTTPException(status_code=400, detail="Inactive user") + return current_user diff --git a/private_gpt/users/api/v1/routers/auth.py b/private_gpt/users/api/v1/routers/auth.py index ba3e4ffe..e3b89a2e 100644 --- a/private_gpt/users/api/v1/routers/auth.py +++ b/private_gpt/users/api/v1/routers/auth.py @@ -169,7 +169,7 @@ def login_access_token( "user": token_payload, "token_type": "bearer", } - log_audit(model='User', action='update', + log_audit(model='User', action='login', details=token_payload, user_id=user.id) return JSONResponse(content=response_dict) @@ -213,7 +213,7 @@ def register( role_name: str = Body(None, title="Role Name", description="User role name (if applicable)"), current_user: models.User = Security( - deps.get_current_user, + deps.get_current_active_user, scopes=[Role.ADMIN["name"], Role.SUPER_ADMIN["name"]], ), ) -> Any: diff --git a/private_gpt/users/api/v1/routers/departments.py b/private_gpt/users/api/v1/routers/departments.py index c8ec0fa4..bab12f1d 100644 --- a/private_gpt/users/api/v1/routers/departments.py +++ b/private_gpt/users/api/v1/routers/departments.py @@ -45,8 +45,31 @@ def list_departments( """ Retrieve a list of departments with pagination support. """ - departments = crud.department.get_multi(db, skip=skip, limit=limit) - return departments + try: + role = current_user.user_role.role.name if current_user.user_role else None + if role == "SUPER_ADMIN": + deps = crud.department.get_multi(db, skip=skip, limit=limit) + else: + deps = crud.department.get_multi_department( + db, department_id=current_user.department_id, skip=skip, limit=limit) + + deps = [ + schemas.Department( + id=dep.id, + name=dep.name, + total_users=dep.total_users, + total_documents=dep.total_documents, + ) + for dep in deps + ] + return deps + except Exception as e: + print(traceback.format_exc()) + logger.error(f"There was an error listing the file(s).") + raise HTTPException( + status_code=500, + detail="Internal Server Error", + ) @router.post("/create", response_model=schemas.Department) diff --git a/private_gpt/users/core/security.py b/private_gpt/users/core/security.py index b053d75f..08c22c2a 100644 --- a/private_gpt/users/core/security.py +++ b/private_gpt/users/core/security.py @@ -7,7 +7,7 @@ from typing import Dict, Any, Optional, Union from jose import JWTError, jwt from passlib.context import CryptContext -ACCESS_TOKEN_EXPIRE_MINUTES = 120 # 30 minutes +ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 12 # 12 hrs REFRESH_TOKEN_EXPIRE_MINUTES = 60 * 24 * 7 # 7 days ALGORITHM = "HS256" # JWT_SECRET_KEY = os.environ['JWT_SECRET_KEY'] # should be kept secret diff --git a/private_gpt/users/crud/department_crud.py b/private_gpt/users/crud/department_crud.py index 3f5586fb..8c6c8376 100644 --- a/private_gpt/users/crud/department_crud.py +++ b/private_gpt/users/crud/department_crud.py @@ -2,7 +2,7 @@ from sqlalchemy.orm import Session from private_gpt.users.schemas.department import DepartmentCreate, DepartmentUpdate from private_gpt.users.models.department import Department from private_gpt.users.crud.base import CRUDBase -from typing import Optional +from typing import Optional, List class CRUDDepartments(CRUDBase[Department, DepartmentCreate, DepartmentUpdate]): @@ -12,4 +12,16 @@ class CRUDDepartments(CRUDBase[Department, DepartmentCreate, DepartmentUpdate]): def get_by_department_name(self, db: Session, *, name: str) -> Optional[Department]: return db.query(self.model).filter(Department.name == name).first() + def get_multi_department( + self, db: Session, *, department_id: int, skip: int = 0, limit: int = 100 + ) -> List[Department]: + return ( + db.query(self.model) + .filter(Department.department_id == department_id) + .offset(skip) + .limit(limit) + .all() + ) + + department = CRUDDepartments(Department) diff --git a/private_gpt/users/models/user.py b/private_gpt/users/models/user.py index 08ccdb0f..b917585b 100644 --- a/private_gpt/users/models/user.py +++ b/private_gpt/users/models/user.py @@ -13,6 +13,7 @@ from sqlalchemy import event, func, select, update from sqlalchemy.orm import relationship from private_gpt.users.db.base_class import Base from private_gpt.users.models.department import Department + class User(Base): """Models a user table""" __tablename__ = "users" @@ -34,6 +35,8 @@ class User(Base): default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, ) + + password_created = Column(DateTime, nullable=True) company_id = Column(Integer, ForeignKey("companies.id"), nullable=True) company = relationship("Company", back_populates="users") @@ -72,3 +75,29 @@ def update_total_users(mapper, connection, target): ) +@event.listens_for(User, 'before_insert') +def set_password_created(mapper, connection, target): + target.password_created = datetime.datetime.utcnow() + connection.execute( + update(User) + .values(password_created=datetime.datetime.utcnow()) + .where(User.id == target.id) + ) + +@event.listens_for(User, 'before_update', propagate=True) +def check_password_expiry(mapper, connection, target): + if target.password_created and ( + datetime.datetime.utcnow() - target.password_created).days > 90: + target.is_active = False + connection.execute( + update(User) + .values(is_active=False) + .where(User.id == target.id) + ) + else: + connection.execute( + update(User) + .values(is_active=True) + .where(User.id == target.id) + ) + diff --git a/private_gpt/users/schemas/__init__.py b/private_gpt/users/schemas/__init__.py index a92b0625..20f3910e 100644 --- a/private_gpt/users/schemas/__init__.py +++ b/private_gpt/users/schemas/__init__.py @@ -1,6 +1,6 @@ from .role import Role, RoleCreate, RoleInDB, RoleUpdate from .token import TokenSchema, TokenPayload -from .user import User, UserCreate, UserInDB, UserUpdate, UserBaseSchema, Profile, UsernameUpdate, DeleteUser, UserAdminUpdate, UserAdmin +from .user import User, UserCreate, UserInDB, UserUpdate, UserBaseSchema, Profile, UsernameUpdate, DeleteUser, UserAdminUpdate, UserAdmin, PasswordUpdate from .user_role import UserRole, UserRoleCreate, UserRoleInDB, UserRoleUpdate from .subscription import Subscription, SubscriptionBase, SubscriptionCreate, SubscriptionUpdate from .company import Company, CompanyBase, CompanyCreate, CompanyUpdate diff --git a/private_gpt/users/schemas/department.py b/private_gpt/users/schemas/department.py index 9e7631a8..baf54d2d 100644 --- a/private_gpt/users/schemas/department.py +++ b/private_gpt/users/schemas/department.py @@ -33,6 +33,10 @@ class DepartmentAdminCreate(DepartmentBase): class Config: orm_mode = True -class Department(DepartmentInDB): - pass +class Department(BaseModel): + id: int + name: str + total_users: Optional[int] + total_documents: Optional[int] + diff --git a/private_gpt/users/schemas/user.py b/private_gpt/users/schemas/user.py index 9544d3b2..8e5c5105 100644 --- a/private_gpt/users/schemas/user.py +++ b/private_gpt/users/schemas/user.py @@ -75,4 +75,8 @@ class UserAdminUpdate(BaseModel): class UserAdmin(BaseModel): fullname: str - department_id: int \ No newline at end of file + department_id: int + + +class PasswordUpdate(BaseModel): + password_created: Optional[datetime] = None