add database more info

Author: ishworgiri1999

.env.template .env.docker.template

@@ -11,4 +11,6 @@ POSTGRES_PASSWORD=postgres
 ENVIRONMENT=development
 
 # Note there maybe be some confusion with urls, use named services instead of
-# localhost inside docker-compose.yml for containers to communicate with each other.
+# localhost inside docker-compose.yml for containers to communicate with each other.
+
+

.gitignore

@@ -9,6 +9,16 @@
 .idea/**/dictionaries
 .idea/**/shelf
 
+# CVE Search plugin data
+
+.cve_search_data/
+
+.DS_Store
+
+.env.docker
+
+cve-search 
+
 # AWS User-specific
 .idea/**/aws.xml
 

Makefile

@@ -1,4 +1,10 @@
 
+.ONESHELL: # Applies to every targets in the file!
+
+all:
+	cd ~/some_dir
+	pwd # Prints ~/some_dir if cd succeeded
+
 
 db-schema:
 ifeq ($(name),)
@@ -17,4 +23,4 @@ db-migrate-down-base:
 
 
 dev:
-	cd src &&  uvicorn app:app --host 0.0.0.0 --port 8001 --reload
+	cd src && uvicorn app:app --host 0.0.0.0 --port 8001 --reload 

docker-compose.yml

@@ -8,7 +8,7 @@ services:
     depends_on:
       - db
     env_file:
-      - .env
+      - .env.docker
     develop:
       watch:
         - action: sync+restart
@@ -56,7 +56,26 @@ services:
       dockerfile: ./dashboard/Dockerfile
     ports:
       - "3000:3000"
+
+
+  # cve_search:
+  #   extends:
+  #     file: cve-search/docker-compose.yml
+  #     service: cve_search
+  #   environment:
+  #     - NVD_NIST_API_KEY=${NVD_NIST_API_KEY}
+  # mongo:
+  #   extends:
+  #     file: cve-search/docker-compose.yml
+  #     service: mongo
+  # redis:
+  #   extends:
+  #     file: cve-search/docker-compose.yml
+  #     service: redis
 volumes:
   ollama:
   db:
   dashboard:
+
+# networks:
+#   cve_search:

requirements.txt

@@ -1,15 +1,14 @@
-python-dotenv~=0.19.1
-pydantic~=2.7.1
+pydantic
 sqlmodel~=0.0.18
 alembic~=1.13.1
 psycopg
 psycopg-binary
-fastapi==0.68.1
+fastapi
 uvicorn==0.15.0
 python-dotenv==0.19.1
 jsonschema==4.22.0
 requests~=2.31.0
 SQLAlchemy~=2.0.30
 tqdm~=4.66.4
 httpx~=0.27.0
-tenacity~=8.3.0
+tenacity~=8.3.0

run

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+source venv/bin/activate
+uvicorn src.app:app --host 0.0.0.0 --port 8000

src/app.py

@@ -1,20 +1,23 @@
+import datetime
 import time
-
-from fastapi import FastAPI, Request
+from typing import Annotated
+from fastapi import Body, FastAPI,Response,Query
 from fastapi.middleware.cors import CORSMiddleware
 
 import api.ollama as ollama
 from data.helper import get_content_list
-from data.db_types import Recommendation, Response
+from data.types import Recommendation, Response,Content
 from my_db import Session
-from models.models import Finding
-from models.models import Recommendation as DBRecommendation
+import models.models as db_models
+
 
-# from data.Findings import Findings
+import data.apischema as apischema
 
+from sqlalchemy import Date, cast
 
 app = FastAPI()
 
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=['*'],
@@ -47,42 +50,51 @@ def health():
 
 
 @app.post('/upload')
-async def upload(request: Request, response: Response):
+async def upload(data: Annotated[apischema.StartRecommendationTaskRequest,Body(...)]):
     """
     This function takes the string from the request and converts it to a data object.
     :return: 200 OK if the data is valid, 400 BAD REQUEST otherwise.
     """
 
-    try:
-        json_data = await request.json()
-    except Exception as e:
-        response.status_code = 400
-        return 'Invalid JSON data'
-
-    # Check if the JSON data is valid
-    # TODO: fix required properties for eg cvss_rating_list is not required
-    # if not validate_json(json_data):
-    #     return 'Invalid JSON data', 400
-
-    # Convert into Response object
-    response = Response.validate(json_data)
-
     # get the content list
-    content_list = get_content_list(response)
+    content_list = get_content_list(data.data)
 
     with Session() as s:
-        for c in content_list:
-            find = Finding(finding=c.title_list[0].element, content=c.json())
-            s.add(find)
-        s.commit()
+            today = datetime.datetime.now().date()
+            existing_task = s.query(db_models.RecommendationTask).filter(cast(db_models.RecommendationTask.created_at,Date)== today).order_by(db_models.RecommendationTask.created_at).first()
+            if existing_task and data.force_update is False:
+                
+                return 'Recommendation task already exists for today', 400
+            
+            if(data.force_update):
+                # Will nuke old task with all its findings.
+                if(existing_task.status == db_models.TaskStatus.PENDING):
+                    return 'Recommendation task is already processing, cannot exit', 400
+                
+                s.query(db_models.RecommendationTask).filter(db_models.RecommendationTask.id == existing_task.id).delete()
+                s.commit()
+                s.delete(existing_task)
+                s.commit()
+            
+            recommendation_task = db_models.RecommendationTask()
+            
+            s.add(recommendation_task)
+            s.commit()
+            s.flush()
+            s.refresh(recommendation_task)
+            for c in content_list:
+                find = db_models.Finding().from_data(c)
+                find.recommendation_task_id = recommendation_task.id
+                s.add(find)
+            s.commit()
     # start subprocess for processing the data
     # ...
 
     return 'Data uploaded successfully'
 
 
 @app.get('/recommendations')
-def recommendations():
+def recommendations(request: Annotated[apischema.GetRecommendationRequest,Query(...)]):
     """
     This function returns the recommendations from the data.
     :return: 200 OK with the recommendations or 204 NO CONTENT if there are no recommendations with retry-after header.
@@ -91,13 +103,20 @@ def recommendations():
     # get the findings
     # ...
     with Session() as s:
-        recs = s.query(DBRecommendation).all()
-
-    recommendations = [Recommendation(recommendation='aa', generic=True) for r in recs]
-    # get the recommendations
-
-    if recommendations:
-        return recommendations, 200
+        total_count = s.query(db_models.Finding).count()
+        recs = s.query(db_models.Finding).join(db_models.RecommendationTask).offset(request.pagination.offset).limit(request.pagination.limit).all()
+        findings = apischema.GetRecommendationResponse(
+            items=[apischema.GetRecommendationResponseItem(
+                description_short="this is a short description",
+                description_long="this is a long description",
+                finding='finding'
+                
+            ) for r in recs],
+        pagination=apischema.Pagination(offset=request.pagination.offset, limit=request.pagination.limit,total=total_count )
+        )
+    
+    if findings:
+        return findings, 200
     else:
         return 'No recommendations found', 204, {'Retry-After': 60}
 

src/data/apischema.py

@@ -0,0 +1,54 @@
+
+
+
+
+from pydantic import BaseModel
+from typing import Optional
+
+from data.pagination import Pagination,PaginationInput
+from data.types import Recommendation, Content,Response
+
+
+class StartRecommendationTaskRequest(BaseModel):
+    user_id: int
+    data: Response
+    force_update: Optional[bool] = False
+    
+class StartRecommendationTaskResponse(BaseModel):
+    task_id: int
+
+
+class GetRecommendationFilter(BaseModel):
+    location:Optional[str]
+    severity: Optional[str]
+    cve_id: Optional[str]
+    source: Optional[str]
+
+
+class GetRecommendationRequest(BaseModel):
+    user_id: int
+    filter:  Optional[GetRecommendationFilter]
+    pagination: Optional[PaginationInput] = PaginationInput(offset=0, limit=10)
+    
+
+    
+class GetRecommendationResponseItem(BaseModel):
+    description_short: str
+    description_long: str
+    finding: str
+    
+class GetRecommendationResponse(BaseModel):
+    items:  list[GetRecommendationResponseItem]
+    pagination: Pagination
+    
+    
+class GetSummarizedRecommendationRequest(BaseModel):
+    user_id: int
+    pagination: PaginationInput = PaginationInput(offset=0, limit=10)
+    
+
+class GetSummarizedRecommendationResponse(BaseModel):
+    recommendation: list[Recommendation]
+    pagination: Pagination
+    
+    

src/data/helper.py

@@ -2,7 +2,7 @@
 import jsonschema
 from jsonschema import validate
 
-from .db_types import schema, Response, Content
+from .types import schema, Response, Content
 
 ##TODO:maybe using pydantic should be enough
 def validate_json(data: any) -> bool:

src/data/db_types.py src/data/types.py

@@ -116,6 +116,10 @@ class Rule(BaseModel):
 
 
 class CveId(BaseModel):
+    element: str
+    source: str
+
+class CweId(BaseModel):
     element: List[str]
     source: str
 
@@ -139,7 +143,8 @@ class Content(BaseModel):
     internal_ratingsource_list: List[Rating]
     cvss_rating_list: Optional[ List[CvssRating]] = None
     rule_list: List[Rule]
-    cwe_id_list: Optional[ List[CveId]] = None
+    cve_id_list: Optional[ List[CveId]] = None
+    cwe_id_list: Optional[ List[CweId]] = None
     activity_list: List[Activity]
     first_found: str
     last_found: str

src/migrations/env.py

@@ -20,7 +20,7 @@
 # target_metadata = None
 
 from models.base import Base
-
+import models
 from models.models import *
 
 target_metadata =  Base.metadata
@@ -36,7 +36,7 @@ def get_url():
     # load env from file
 
 
-    from db import (get_db_url)
+    from my_db import (get_db_url)
     return get_db_url()
 
 

src/migrations/script.py.mako

@@ -8,6 +8,7 @@ Create Date: ${create_date}
 from alembic import op
 import sqlalchemy as sa
 import sqlmodel.sql.sqltypes
+import models
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.

src/migrations/versions/c3fdb30822c4_add_more_tables2.py

@@ -0,0 +1,42 @@
+"""add_more_tables2
+
+Revision ID: c3fdb30822c4
+Revises: cefb826a201b
+Create Date: 2024-06-11 17:59:48.102377
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+import models
+
+
+# revision identifiers, used by Alembic.
+revision = 'c3fdb30822c4'
+down_revision = 'cefb826a201b'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('findings', 'recommendation_task_id',
+               existing_type=sa.INTEGER(),
+               nullable=False)
+    op.drop_constraint('findings_recommendation_task_id_fkey', 'findings', type_='foreignkey')
+    op.create_foreign_key(None, 'findings', 'recommendation_task', ['recommendation_task_id'], ['id'], ondelete='CASCADE')
+    op.add_column('recommendations', sa.Column('recommendation_task_id', sa.Integer(), nullable=False))
+    op.create_foreign_key(None, 'recommendations', 'recommendation_task', ['recommendation_task_id'], ['id'], ondelete='CASCADE')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, 'recommendations', type_='foreignkey')
+    op.drop_column('recommendations', 'recommendation_task_id')
+    op.drop_constraint(None, 'findings', type_='foreignkey')
+    op.create_foreign_key('findings_recommendation_task_id_fkey', 'findings', 'recommendation_task', ['recommendation_task_id'], ['id'])
+    op.alter_column('findings', 'recommendation_task_id',
+               existing_type=sa.INTEGER(),
+               nullable=True)
+    # ### end Alembic commands ###