diff --git a/.github/workflows/django-test.yml b/.github/workflows/django-test.yml new file mode 100644 index 00000000..02f22059 --- /dev/null +++ b/.github/workflows/django-test.yml @@ -0,0 +1,41 @@ +name: Django CI + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + +jobs: + build: + + runs-on: ubuntu-latest + strategy: + max-parallel: 4 + matrix: + python-version: ["3.10", "3.11"] + + steps: + - uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v3 + with: + python-version: ${{ matrix.python-version }} + - name: Install Dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pylint + pip install pylint-django + pip install isort + - name: Generate Fernet Key + run: | + cd chirps; + python -c "from cryptography.fernet import Fernet; print(f'FERNET_KEY={Fernet.generate_key().decode()}')" > .env + - name: Run Tests + run: | + cd chirps; python manage.py test + - name: Run Linting + run: | + isort --check-only --diff . + pylint --load-plugins pylint_django --django-settings-module="chirps.settings" $(find . -name "*.py" | xargs) diff --git a/.pylintrc b/.pylintrc index 7615b869..a83ed107 100644 --- a/.pylintrc +++ b/.pylintrc @@ -1,2 +1,10 @@ [FORMAT] max-line-length=120 + +[MASTER] +# Ignore Django migration files +ignore-patterns=\d{4}_.*?.py + +[MESSAGES CONTROL] +disable=duplicate-code,too-few-public-methods,imported-auth-user +ignored-modules=pinecone diff --git a/chirps/account/admin.py b/chirps/account/admin.py index 8c38f3f3..7164d8bc 100644 --- a/chirps/account/admin.py +++ b/chirps/account/admin.py @@ -1,3 +1,24 @@ +"""Define admin interface models for the account application.""" from django.contrib import admin +from django.contrib.auth.admin import UserAdmin as BaseUserAdmin +from django.contrib.auth.models import User -# Register your models here. +from .models import Profile + + +class ProfileInline(admin.StackedInline): + """Inline admin descriptor for the Profile model.""" + + model = Profile + can_delete = False + verbose_name_plural = 'profile' + + +class UserAdmin(BaseUserAdmin): + """Define a new User admin.""" + inlines = [ProfileInline] + + +# Re-register UserAdmin +admin.site.unregister(User) +admin.site.register(User, UserAdmin) diff --git a/chirps/account/apps.py b/chirps/account/apps.py index 2b08f1ad..e5c996b0 100644 --- a/chirps/account/apps.py +++ b/chirps/account/apps.py @@ -1,6 +1,9 @@ +"""Configures the account app.""""" from django.apps import AppConfig class AccountConfig(AppConfig): + """Configuration options for the account application.""" + default_auto_field = 'django.db.models.BigAutoField' name = 'account' diff --git a/chirps/account/forms.py b/chirps/account/forms.py index 332d6906..41609878 100644 --- a/chirps/account/forms.py +++ b/chirps/account/forms.py @@ -1,3 +1,4 @@ +"""Form classes for the account app.""" from django import forms from django.contrib.auth.hashers import make_password from django.forms import ModelForm @@ -6,21 +7,29 @@ class ProfileForm(ModelForm): + """Form for the Profile model.""" + def clean_openai_key(self): + """Hash the openai_key before saving it to the database.""" data = self.cleaned_data['openai_key'] return make_password(data) class Meta: + """Meta class for ProfileForm.""" model = Profile fields = ['openai_key'] widgets = {'openai_key': forms.PasswordInput(attrs={'class': 'form-control'})} class LoginForm(forms.Form): + """Form for logging in.""" + username = forms.CharField(max_length=256) password = forms.CharField(max_length=256, widget=forms.PasswordInput) class SignupForm(forms.Form): + """Form for signing up.""" + username = forms.CharField(max_length=256) email = forms.EmailField(max_length=256) password1 = forms.CharField(max_length=256, widget=forms.PasswordInput) diff --git a/chirps/account/migrations/0001_initial.py b/chirps/account/migrations/0001_initial.py index 70ea1ee0..e278a95f 100644 --- a/chirps/account/migrations/0001_initial.py +++ b/chirps/account/migrations/0001_initial.py @@ -1,8 +1,8 @@ # Generated by Django 4.2.2 on 2023-06-29 13:56 +import django.db.models.deletion from django.conf import settings from django.db import migrations, models -import django.db.models.deletion class Migration(migrations.Migration): diff --git a/chirps/account/models.py b/chirps/account/models.py index 939494f5..b6c0a419 100644 --- a/chirps/account/models.py +++ b/chirps/account/models.py @@ -1,27 +1,10 @@ -from django.contrib import admin -from django.contrib.auth.admin import UserAdmin as BaseUserAdmin +"""Models for the account appliation.""" from django.contrib.auth.models import User from django.db import models class Profile(models.Model): + """Custom profile model for users.""" + user = models.OneToOneField(User, on_delete=models.CASCADE) openai_key = models.CharField(max_length=100, blank=True) - - -# Define an inline admin descriptor for Employee model -# which acts a bit like a singleton -class ProfileInline(admin.StackedInline): - model = Profile - can_delete = False - verbose_name_plural = 'profile' - - -# Define a new User admin -class UserAdmin(BaseUserAdmin): - inlines = [ProfileInline] - - -# Re-register UserAdmin -admin.site.unregister(User) -admin.site.register(User, UserAdmin) diff --git a/chirps/account/tests.py b/chirps/account/tests.py index 8abf3ea6..c141b067 100644 --- a/chirps/account/tests.py +++ b/chirps/account/tests.py @@ -1,10 +1,13 @@ -from django.contrib.auth.hashers import check_password, make_password +"""Tests for the account application.""" from django.test import TestCase from django.urls import reverse + from .forms import ProfileForm -from django.contrib.auth.forms import AuthenticationForm + class AccountTests(TestCase): + """Main test class for the account application.""" + def test_openai_key_hash(self): """Verify that the openai_key paramater is correctly hashed by the form""" secret_val = 'secret_12345abcd' diff --git a/chirps/account/urls.py b/chirps/account/urls.py index b3cb0f2f..ca8dd0e7 100644 --- a/chirps/account/urls.py +++ b/chirps/account/urls.py @@ -1,11 +1,12 @@ +"""URLs for the account app.""" from django.contrib.auth import views as auth_views from django.urls import path +from .views import login_view from .views import profile as profile_view -from .views import signup, login_view +from .views import signup urlpatterns = [ - # path('login/', auth_views.LoginView.as_view(template_name='account/login.html', next_page='/'), name='login'), path('login/', login_view, name='login'), path( 'logout/', auth_views.LogoutView.as_view(template_name='account/logout.html', next_page='login'), name='logout' diff --git a/chirps/account/views.py b/chirps/account/views.py index 70be99e2..4743fd65 100644 --- a/chirps/account/views.py +++ b/chirps/account/views.py @@ -1,22 +1,23 @@ +"""Views for the account application.""" +from django.contrib.auth import authenticate, login +from django.contrib.auth.models import User # noqa: E5142 from django.shortcuts import redirect, render from django.urls import reverse -from django.contrib.auth.models import User -from django.contrib.auth import login + +from .forms import LoginForm, ProfileForm, SignupForm from .models import Profile -from .forms import ProfileForm, SignupForm, LoginForm -from django.contrib.auth import authenticate def profile(request): - + """Render the user profile page and handle updates""" if request.method == 'POST': form = ProfileForm(request.POST, instance=request.user.profile) if form.is_valid(): - profile = form.save(commit=False) - profile.user = request.user - profile.save() + profile_form = form.save(commit=False) + profile_form.user = request.user + profile_form.save() # Redirect the user back to the dashboard return redirect('profile') @@ -27,6 +28,7 @@ def profile(request): return render(request, 'account/profile.html', {'form': form}) def signup(request): + """Render the signup page and handle posts.""" if request.method == 'POST': form = SignupForm(request.POST) if form.is_valid(): @@ -53,8 +55,8 @@ def signup(request): user.save() # Create the user profile - profile = Profile(user=user) - profile.save() + user_profile = Profile(user=user) + user_profile.save() # Login the user login(request, user) @@ -67,6 +69,7 @@ def signup(request): return render(request, 'account/signup.html', {'form': form}) def login_view(request): + """Render the login page.""" # If there are no users, redirect to the installation page if User.objects.count() == 0: diff --git a/chirps/base_app/admin.py b/chirps/base_app/admin.py index 8c38f3f3..512eaba7 100644 --- a/chirps/base_app/admin.py +++ b/chirps/base_app/admin.py @@ -1,3 +1 @@ -from django.contrib import admin - -# Register your models here. +"""Define any admin interface models for the base application.""" diff --git a/chirps/base_app/apps.py b/chirps/base_app/apps.py index 4b2d8549..95ac10a7 100644 --- a/chirps/base_app/apps.py +++ b/chirps/base_app/apps.py @@ -1,6 +1,9 @@ +"""Application configration for the base_app app.""" from django.apps import AppConfig class BaseAppConfig(AppConfig): + """Application configration for the base_app app.""" + default_auto_field = 'django.db.models.BigAutoField' name = 'base_app' diff --git a/chirps/base_app/forms.py b/chirps/base_app/forms.py index de0030b1..e5fa8155 100644 --- a/chirps/base_app/forms.py +++ b/chirps/base_app/forms.py @@ -1,9 +1,14 @@ +"""Forms for the base application.""" from django import forms -class InstallForm(forms.Form): +class InstallForm(forms.Form): + """Form to render the new installation page.""" superuser_username = forms.CharField(label='Superuser Username', max_length=100) superuser_email = forms.EmailField(label='Superuser Email', max_length=100) superuser_password = forms.CharField(label='Superuser Password', max_length=100, widget=forms.PasswordInput) - superuser_password_confirm = forms.CharField(label='Superuser Password (Confirm)', max_length=100, widget=forms.PasswordInput) - + superuser_password_confirm = forms.CharField( + label='Superuser Password (Confirm)', + max_length=100, + widget=forms.PasswordInput + ) diff --git a/chirps/base_app/management/commands/celery.py b/chirps/base_app/management/commands/celery.py index e6606f36..ee2c040f 100644 --- a/chirps/base_app/management/commands/celery.py +++ b/chirps/base_app/management/commands/celery.py @@ -1,9 +1,11 @@ +"""Celery management command.""" import os from django.core.management.base import BaseCommand class Command(BaseCommand): + """Manage a local celery installation with this command.""" help = 'Interact with the local celery broker' def add_arguments(self, parser): @@ -13,7 +15,7 @@ def add_arguments(self, parser): parser.add_argument('--restart', action='store_true', help='Restart celery server') def handle(self, *args, **options): - + """Handle the command.""" if options['start']: self.start() elif options['stop']: @@ -23,9 +25,11 @@ def handle(self, *args, **options): self.start() def start(self): + """Start the celery server.""" os.system('sudo mkdir -p /var/run/celery; sudo chmod 777 /var/run/celery') os.system('sudo mkdir -p /var/log/celery; sudo chmod 777 /var/log/celery') os.system('celery multi start w1 -A chirps -l INFO') def stop(self): + """Stop the celery server.""" os.system('celery multi stopwait w1 -A chirps -l INFO') diff --git a/chirps/base_app/management/commands/initialize_app.py b/chirps/base_app/management/commands/initialize_app.py index 77ebd55e..6d7b38b4 100644 --- a/chirps/base_app/management/commands/initialize_app.py +++ b/chirps/base_app/management/commands/initialize_app.py @@ -1,56 +1,60 @@ +"""Management command to initialize the app by running multiple management commands in succession.""" import os + from django.contrib.auth.models import User -from django.core.management.base import BaseCommand -from django.core.management import call_command - -class Command(BaseCommand): - help = 'Initialize the app by running multiple management commands' - - def handle(self, *args, **options): - # Run the 'redis --start' command - self.stdout.write(self.style.WARNING('Starting Redis...')) - call_command('redis', '--start') - self.stdout.write(self.style.SUCCESS('Redis started')) - - # Run the 'rabbitmq --start' command - self.stdout.write(self.style.WARNING('Starting RabbitMQ...')) - call_command('rabbitmq', '--start') - self.stdout.write(self.style.SUCCESS('RabbitMQ started')) - - # Run the 'celery --start' command - self.stdout.write(self.style.WARNING('Starting Celery...')) +from django.core.management import call_command +from django.core.management.base import BaseCommand + + +class Command(BaseCommand): + """Initialize the app by running multiple management commands.""" + help = 'Initialize the app by running multiple management commands' + + def handle(self, *args, **options): + # Run the 'redis --start' command + self.stdout.write(self.style.WARNING('Starting Redis...')) + call_command('redis', '--start') + self.stdout.write(self.style.SUCCESS('Redis started')) + + # Run the 'rabbitmq --start' command + self.stdout.write(self.style.WARNING('Starting RabbitMQ...')) + call_command('rabbitmq', '--start') + self.stdout.write(self.style.SUCCESS('RabbitMQ started')) + + # Run the 'celery --start' command + self.stdout.write(self.style.WARNING('Starting Celery...')) os.system('sudo mkdir -p /var/run/celery; sudo chmod 777 /var/run/celery') os.system('sudo mkdir -p /var/log/celery; sudo chmod 777 /var/log/celery') os.system('celery multi start w1 -A chirps -l INFO') - self.stdout.write(self.style.SUCCESS('Celery started')) - - # Run the 'makemigrations' command - self.stdout.write(self.style.WARNING('Running makemigrations...')) - call_command('makemigrations') - self.stdout.write(self.style.SUCCESS('makemigrations completed')) - - # Run the 'migrate' command - self.stdout.write(self.style.WARNING('Running migrate...')) - call_command('migrate') - self.stdout.write(self.style.SUCCESS('migrate completed')) - - # Check if a superuser already exists - if not User.objects.filter(is_superuser=True).exists(): - # Run the 'createsuperuser' command - self.stdout.write(self.style.WARNING('Creating superuser...')) - call_command('createsuperuser') - self.stdout.write(self.style.SUCCESS('Superuser created')) - else: - self.stdout.write(self.style.WARNING('Superuser already exists. Skipping superuser creation.')) - - # Run the 'loaddata' command - self.stdout.write(self.style.WARNING('Loading data from fixtures...')) - call_command('loaddata', 'plan/fixtures/plan/employee.json') - self.stdout.write(self.style.SUCCESS('Data loaded from fixtures')) - - # Run the 'runserver' command - self.stdout.write(self.style.WARNING('Starting the development server...')) - call_command('runserver') - self.stdout.write(self.style.SUCCESS('Development server started')) - - self.stdout.write(self.style.SUCCESS('App initialization completed')) + self.stdout.write(self.style.SUCCESS('Celery started')) + + # Run the 'makemigrations' command + self.stdout.write(self.style.WARNING('Running makemigrations...')) + call_command('makemigrations') + self.stdout.write(self.style.SUCCESS('makemigrations completed')) + + # Run the 'migrate' command + self.stdout.write(self.style.WARNING('Running migrate...')) + call_command('migrate') + self.stdout.write(self.style.SUCCESS('migrate completed')) + + # Check if a superuser already exists + if not User.objects.filter(is_superuser=True).exists(): + # Run the 'createsuperuser' command + self.stdout.write(self.style.WARNING('Creating superuser...')) + call_command('createsuperuser') + self.stdout.write(self.style.SUCCESS('Superuser created')) + else: + self.stdout.write(self.style.WARNING('Superuser already exists. Skipping superuser creation.')) + + # Run the 'loaddata' command + self.stdout.write(self.style.WARNING('Loading data from fixtures...')) + call_command('loaddata', 'plan/fixtures/plan/employee.json') + self.stdout.write(self.style.SUCCESS('Data loaded from fixtures')) + + # Run the 'runserver' command + self.stdout.write(self.style.WARNING('Starting the development server...')) + call_command('runserver') + self.stdout.write(self.style.SUCCESS('Development server started')) + + self.stdout.write(self.style.SUCCESS('App initialization completed')) diff --git a/chirps/base_app/management/commands/rabbitmq.py b/chirps/base_app/management/commands/rabbitmq.py index 5f7d06a4..4b8dc96a 100644 --- a/chirps/base_app/management/commands/rabbitmq.py +++ b/chirps/base_app/management/commands/rabbitmq.py @@ -1,9 +1,11 @@ +"""Management command for interacting with rabbitmq.""" import os from django.core.management.base import BaseCommand class Command(BaseCommand): + """Management command for interacting with rabbitmq.""" help = 'Interact with the local rabbitmq development server' def add_arguments(self, parser): diff --git a/chirps/base_app/management/commands/redis.py b/chirps/base_app/management/commands/redis.py index b560ce1c..d098ba01 100644 --- a/chirps/base_app/management/commands/redis.py +++ b/chirps/base_app/management/commands/redis.py @@ -1,9 +1,11 @@ +"""Management command for interacting with redis.""" import os from django.core.management.base import BaseCommand class Command(BaseCommand): + """Management command for interacting with redis.""" help = 'Interact with the local redis development server' def add_arguments(self, parser): diff --git a/chirps/base_app/management/commands/start_services.py b/chirps/base_app/management/commands/start_services.py index 79626fdc..28303a17 100644 --- a/chirps/base_app/management/commands/start_services.py +++ b/chirps/base_app/management/commands/start_services.py @@ -1,32 +1,35 @@ +"""Management command to start the app services.""" import os -from django.contrib.auth.models import User -from django.core.management.base import BaseCommand -from django.core.management import call_command - -class Command(BaseCommand): - help = 'Initialize the app by running multiple management commands' - - def handle(self, *args, **options): - # Run the 'redis --start' command - self.stdout.write(self.style.WARNING('Starting Redis...')) - call_command('redis', '--start') - self.stdout.write(self.style.SUCCESS('Redis started')) - - # Run the 'rabbitmq --start' command - self.stdout.write(self.style.WARNING('Starting RabbitMQ...')) - call_command('rabbitmq', '--start') - self.stdout.write(self.style.SUCCESS('RabbitMQ started')) - - # Run the 'celery --start' command - self.stdout.write(self.style.WARNING('Starting Celery...')) + +from django.core.management import call_command +from django.core.management.base import BaseCommand + + +class Command(BaseCommand): + """Initialize the app by running multiple management commands.""" + help = 'Initialize the app by running multiple management commands' + + def handle(self, *args, **options): + # Run the 'redis --start' command + self.stdout.write(self.style.WARNING('Starting Redis...')) + call_command('redis', '--start') + self.stdout.write(self.style.SUCCESS('Redis started')) + + # Run the 'rabbitmq --start' command + self.stdout.write(self.style.WARNING('Starting RabbitMQ...')) + call_command('rabbitmq', '--start') + self.stdout.write(self.style.SUCCESS('RabbitMQ started')) + + # Run the 'celery --start' command + self.stdout.write(self.style.WARNING('Starting Celery...')) os.system('sudo mkdir -p /var/run/celery; sudo chmod 777 /var/run/celery') os.system('sudo mkdir -p /var/log/celery; sudo chmod 777 /var/log/celery') os.system('celery multi start w1 -A chirps -l INFO') self.stdout.write(self.style.SUCCESS('Celery started')) - # Run the 'runserver' command - self.stdout.write(self.style.WARNING('Starting the development server...')) - call_command('runserver') - self.stdout.write(self.style.SUCCESS('Development server started')) - - self.stdout.write(self.style.SUCCESS('App initialization completed')) \ No newline at end of file + # Run the 'runserver' command + self.stdout.write(self.style.WARNING('Starting the development server...')) + call_command('runserver') + self.stdout.write(self.style.SUCCESS('Development server started')) + + self.stdout.write(self.style.SUCCESS('App initialization completed')) diff --git a/chirps/base_app/models.py b/chirps/base_app/models.py index 71a83623..2c5bd3bc 100644 --- a/chirps/base_app/models.py +++ b/chirps/base_app/models.py @@ -1,3 +1 @@ -from django.db import models - -# Create your models here. +"""Define any models for the chirps base application.""" diff --git a/chirps/base_app/tests.py b/chirps/base_app/tests.py index 7ce503c2..df344c35 100644 --- a/chirps/base_app/tests.py +++ b/chirps/base_app/tests.py @@ -1,3 +1 @@ -from django.test import TestCase - -# Create your tests here. +"""Tests for the base application.""" diff --git a/chirps/base_app/urls.py b/chirps/base_app/urls.py index 5ea4195e..edb5c1a6 100644 --- a/chirps/base_app/urls.py +++ b/chirps/base_app/urls.py @@ -1,3 +1,4 @@ +"""Top level URLS for the project""" from django.urls import path from . import views diff --git a/chirps/base_app/views.py b/chirps/base_app/views.py index 3f1d389e..28839cd5 100644 --- a/chirps/base_app/views.py +++ b/chirps/base_app/views.py @@ -1,18 +1,21 @@ -from django.contrib.auth.decorators import login_required -from django.shortcuts import render, redirect - -from django.contrib.auth.models import User +"""Views for the base application.""" +from account.models import Profile from django.contrib.auth import login +from django.contrib.auth.decorators import login_required +from django.contrib.auth.models import User # noqa: E5142 +from django.shortcuts import redirect, render from django.urls import reverse -from . forms import InstallForm -from account.models import Profile +from .forms import InstallForm + @login_required def index(request): + """Render the index page.""" return render(request, 'dashboard/index.html', {}) def install(request): + """Render the install page.""" # If there are uses already, redirect to the dashboard if User.objects.count() > 0: diff --git a/chirps/chirps/__init__.py b/chirps/chirps/__init__.py index 15d7c508..4a1e8a4f 100644 --- a/chirps/chirps/__init__.py +++ b/chirps/chirps/__init__.py @@ -1,5 +1,4 @@ -# This will make sure the app is always imported when -# Django starts so that shared_task will use this app. +"""Make sure the app is always imported when Django starts so that shared_task will use this app.""" from .celery import app as celery_app __all__ = ('celery_app',) diff --git a/chirps/chirps/celery.py b/chirps/chirps/celery.py index d26466f2..4e0e6110 100644 --- a/chirps/chirps/celery.py +++ b/chirps/chirps/celery.py @@ -1,3 +1,5 @@ +"""Celery configuration for the chirps project.""" + import os from celery import Celery @@ -15,8 +17,3 @@ # Load task modules from all registered Django apps. app.autodiscover_tasks() - - -@app.task(bind=True, ignore_result=True) -def debug_task(self): - print(f'Request: {self.request!r}') diff --git a/chirps/chirps/key_script.py b/chirps/chirps/key_script.py index 649c41d4..ced03864 100644 --- a/chirps/chirps/key_script.py +++ b/chirps/chirps/key_script.py @@ -1,4 +1,5 @@ -from cryptography.fernet import Fernet - -key = Fernet.generate_key() -print(key.decode()) +"""Helper script to generate a fernet key.""" +from cryptography.fernet import Fernet + +key = Fernet.generate_key() +print(key.decode()) diff --git a/chirps/chirps/settings.py b/chirps/chirps/settings.py index 5b1adf26..30d28c88 100644 --- a/chirps/chirps/settings.py +++ b/chirps/chirps/settings.py @@ -10,23 +10,21 @@ https://docs.djangoproject.com/en/4.2/ref/settings/ """ import os -from dotenv import load_dotenv - -load_dotenv(os.path.join(os.path.dirname(os.path.dirname(__file__)), '.env')) - from pathlib import Path -# HACK: (alexn) monkeypatching because django 4.0 does not have force_text -# see https://stackoverflow.com/a/70833150 import django from django.utils.encoding import force_str +from dotenv import load_dotenv + +# HACK: (alexn) monkeypatching because django 4.0 does not have force_text +# see https://stackoverflow.com/a/70833150 django.utils.encoding.force_text = force_str +load_dotenv(os.path.join(os.path.dirname(os.path.dirname(__file__)), '.env')) # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = Path(__file__).resolve().parent.parent - # Quick-start development settings - unsuitable for production # See https://docs.djangoproject.com/en/4.2/howto/deployment/checklist/ @@ -149,4 +147,7 @@ CELERY_CACHE_BACKEND = 'django-cache' # FERNET SETTINGS -FERNET_KEY = os.getenv('FERNET_KEY') \ No newline at end of file +if os.getenv('FERNET_KEY') is None: + raise Exception('FERNET_KEY environment variable is not set') # pylint: disable=broad-exception-raised + +FERNET_KEY = os.getenv('FERNET_KEY') diff --git a/chirps/manage.py b/chirps/manage.py index d038b06f..4b8dd104 100755 --- a/chirps/manage.py +++ b/chirps/manage.py @@ -8,7 +8,8 @@ def main(): """Run administrative tasks.""" os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'chirps.settings') try: - from django.core.management import execute_from_command_line + from django.core.management import \ + execute_from_command_line # pylint: disable=import-outside-toplevel except ImportError as exc: raise ImportError( "Couldn't import Django. Are you sure it's installed and " diff --git a/chirps/plan/admin.py b/chirps/plan/admin.py index 8c38f3f3..5aa22d7a 100644 --- a/chirps/plan/admin.py +++ b/chirps/plan/admin.py @@ -1,3 +1,7 @@ +"""Admin interface definition for plan models.""" from django.contrib import admin -# Register your models here. +from .models import Plan, Rule + +admin.site.register(Plan) +admin.site.register(Rule) diff --git a/chirps/plan/apps.py b/chirps/plan/apps.py index 451b518e..74eec87b 100644 --- a/chirps/plan/apps.py +++ b/chirps/plan/apps.py @@ -1,6 +1,8 @@ +"""Configuration options for the plan application.""" from django.apps import AppConfig class PlanConfig(AppConfig): + """Configuration options for the plan application.""" default_auto_field = 'django.db.models.BigAutoField' name = 'plan' diff --git a/chirps/plan/management/__init__.py b/chirps/plan/management/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/chirps/plan/management/commands/__init__.py b/chirps/plan/management/commands/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/chirps/plan/management/commands/create_templates.py b/chirps/plan/management/commands/create_templates.py deleted file mode 100644 index aca3e165..00000000 --- a/chirps/plan/management/commands/create_templates.py +++ /dev/null @@ -1,23 +0,0 @@ -import os - -from django.core.management.base import BaseCommand - - -class Command(BaseCommand): - help = 'Interact with the local celery broker' - - def add_arguments(self, parser): - - parser.add_argument('--start', action='store_true', help='Start celery server') - parser.add_argument('--stop', action='store_true', help='Stop celery server') - parser.add_argument('--restart', action='store_true', help='Restart celery server') - - def handle(self, *args, **options): - - if options['start']: - self.start() - elif options['stop']: - self.stop() - elif options['restart']: - self.stop() - self.start() diff --git a/chirps/plan/migrations/0001_initial.py b/chirps/plan/migrations/0001_initial.py index da5a86e3..e4a67968 100644 --- a/chirps/plan/migrations/0001_initial.py +++ b/chirps/plan/migrations/0001_initial.py @@ -1,8 +1,8 @@ # Generated by Django 4.2.2 on 2023-06-29 13:56 +import django.db.models.deletion from django.conf import settings from django.db import migrations, models -import django.db.models.deletion class Migration(migrations.Migration): diff --git a/chirps/plan/models.py b/chirps/plan/models.py index 7e208eb6..58228a03 100644 --- a/chirps/plan/models.py +++ b/chirps/plan/models.py @@ -1,4 +1,5 @@ -from django.contrib import admin +"""Models for the plan application.""" +from django.contrib.auth.models import User from django.db import models @@ -12,8 +13,8 @@ class Plan(models.Model): is_template = models.BooleanField(default=False) # Bind this plan to a user if it isn't a template - user = models.ForeignKey('auth.User', on_delete=models.CASCADE, null=True, blank=True) - + user = models.ForeignKey(User, on_delete=models.CASCADE, null=True, blank=True) + def __str__(self): return self.name @@ -41,7 +42,3 @@ class Rule(models.Model): def __str__(self): return self.name - - -admin.site.register(Plan) -admin.site.register(Rule) diff --git a/chirps/plan/tests.py b/chirps/plan/tests.py index 7ce503c2..60f1768e 100644 --- a/chirps/plan/tests.py +++ b/chirps/plan/tests.py @@ -1,3 +1 @@ -from django.test import TestCase - -# Create your tests here. +"""Tests for the plan application.""" diff --git a/chirps/plan/urls.py b/chirps/plan/urls.py index a6036594..911387fa 100644 --- a/chirps/plan/urls.py +++ b/chirps/plan/urls.py @@ -1,3 +1,4 @@ +"""URLs for the plan app.""" from django.urls import path from . import views diff --git a/chirps/plan/views.py b/chirps/plan/views.py index c59d99ff..fa282837 100644 --- a/chirps/plan/views.py +++ b/chirps/plan/views.py @@ -1,3 +1,4 @@ +"""Views for the plan app.""" from django.shortcuts import render from .models import Plan diff --git a/chirps/scan/admin.py b/chirps/scan/admin.py index 8c38f3f3..039245fa 100644 --- a/chirps/scan/admin.py +++ b/chirps/scan/admin.py @@ -1,3 +1,8 @@ +"""Registration point of the admin interface for the scan app.""" from django.contrib import admin -# Register your models here. +from .models import Finding, Result, Scan + +admin.site.register(Scan) +admin.site.register(Result) +admin.site.register(Finding) diff --git a/chirps/scan/apps.py b/chirps/scan/apps.py index a4f8e921..e00c9370 100644 --- a/chirps/scan/apps.py +++ b/chirps/scan/apps.py @@ -1,6 +1,8 @@ +"""Configuration options for the scan app.""" from django.apps import AppConfig class ScanConfig(AppConfig): + """Configuration options for the scan application.""" default_auto_field = 'django.db.models.BigAutoField' name = 'scan' diff --git a/chirps/scan/forms.py b/chirps/scan/forms.py index e8c7e4f7..cee98c4c 100644 --- a/chirps/scan/forms.py +++ b/chirps/scan/forms.py @@ -1,12 +1,14 @@ +"""Forms for rendering scan application models.""" from django import forms from django.forms import ModelForm -from target.models import BaseTarget from .models import Scan class ScanForm(ModelForm): + """Form for the main scan model.""" class Meta: + """Django Meta options for the ScanForm.""" model = Scan fields = ['description', 'target', 'plan'] diff --git a/chirps/scan/migrations/0001_initial.py b/chirps/scan/migrations/0001_initial.py index ad963e18..5f56606d 100644 --- a/chirps/scan/migrations/0001_initial.py +++ b/chirps/scan/migrations/0001_initial.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.2 on 2023-06-29 13:56 -from django.db import migrations, models import django.db.models.deletion +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/chirps/scan/migrations/0002_initial.py b/chirps/scan/migrations/0002_initial.py index e4a21af4..c4d0d7b3 100644 --- a/chirps/scan/migrations/0002_initial.py +++ b/chirps/scan/migrations/0002_initial.py @@ -1,8 +1,8 @@ # Generated by Django 4.2.2 on 2023-06-29 13:56 +import django.db.models.deletion from django.conf import settings from django.db import migrations, models -import django.db.models.deletion class Migration(migrations.Migration): diff --git a/chirps/scan/migrations/0004_alter_result_findings.py b/chirps/scan/migrations/0004_alter_result_findings.py index 478eee61..33bd25bb 100644 --- a/chirps/scan/migrations/0004_alter_result_findings.py +++ b/chirps/scan/migrations/0004_alter_result_findings.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.2 on 2023-06-29 16:29 -from django.db import migrations import fernet_fields.fields +from django.db import migrations class Migration(migrations.Migration): diff --git a/chirps/scan/migrations/0005_result_text.py b/chirps/scan/migrations/0005_result_text.py index b4ae2dc3..d1b826ac 100644 --- a/chirps/scan/migrations/0005_result_text.py +++ b/chirps/scan/migrations/0005_result_text.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.2 on 2023-07-03 18:46 -from django.db import migrations import fernet_fields.fields +from django.db import migrations class Migration(migrations.Migration): diff --git a/chirps/scan/migrations/0007_remove_result_count_remove_result_findings_finding.py b/chirps/scan/migrations/0007_remove_result_count_remove_result_findings_finding.py index e30ef5a6..a591ea00 100644 --- a/chirps/scan/migrations/0007_remove_result_count_remove_result_findings_finding.py +++ b/chirps/scan/migrations/0007_remove_result_count_remove_result_findings_finding.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.2 on 2023-07-03 20:14 -from django.db import migrations, models import django.db.models.deletion +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/chirps/scan/migrations/0009_remove_scan_results_result_scan.py b/chirps/scan/migrations/0009_remove_scan_results_result_scan.py index d58cdb33..20e5332f 100644 --- a/chirps/scan/migrations/0009_remove_scan_results_result_scan.py +++ b/chirps/scan/migrations/0009_remove_scan_results_result_scan.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.2 on 2023-07-03 20:28 -from django.db import migrations, models import django.db.models.deletion +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/chirps/scan/migrations/0010_alter_result_scan.py b/chirps/scan/migrations/0010_alter_result_scan.py index 69da6f83..cbc2f4b1 100644 --- a/chirps/scan/migrations/0010_alter_result_scan.py +++ b/chirps/scan/migrations/0010_alter_result_scan.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.2 on 2023-07-03 20:38 -from django.db import migrations, models import django.db.models.deletion +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/chirps/scan/migrations/0011_alter_result_scan.py b/chirps/scan/migrations/0011_alter_result_scan.py index 6acbc33d..c7378acd 100644 --- a/chirps/scan/migrations/0011_alter_result_scan.py +++ b/chirps/scan/migrations/0011_alter_result_scan.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.2 on 2023-07-03 20:44 -from django.db import migrations, models import django.db.models.deletion +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/chirps/scan/models.py b/chirps/scan/models.py index f5d19b10..5150b0da 100644 --- a/chirps/scan/models.py +++ b/chirps/scan/models.py @@ -1,11 +1,11 @@ """Models for the scan application.""" -from django.contrib import admin -from django.db import models -from django_celery_results.models import TaskResult -from plan.models import Rule from django.contrib.auth.models import User +from django.db import models from django.utils.safestring import mark_safe +from django_celery_results.models import TaskResult from fernet_fields import EncryptedTextField +from plan.models import Rule + class Scan(models.Model): """Model for a single scan run against a target.""" @@ -86,7 +86,3 @@ def with_highlight(self): buffer += "" buffer += self.result.text[self.offset + self.length + 1 : ] return mark_safe(buffer) - -admin.site.register(Scan) -admin.site.register(Result) -admin.site.register(Finding) diff --git a/chirps/scan/tasks.py b/chirps/scan/tasks.py index ce3bdcdc..e1935028 100644 --- a/chirps/scan/tasks.py +++ b/chirps/scan/tasks.py @@ -1,19 +1,16 @@ -import json +"""Celery tasks for the scan application.""" import re + from celery import shared_task from django.utils import timezone from target.models import BaseTarget -from .models import Result, Rule, Scan, Finding - - -@shared_task -def add(x, y): - return x + y +from .models import Finding, Result, Scan @shared_task def scan_task(scan_id): + """Main scan task.""" print(f'Running a scan {scan_id}') try: @@ -32,7 +29,6 @@ def scan_task(scan_id): for rule in scan.plan.rules.all(): print(f'Running rule {rule}') - # TODO: Convert the query to an embedding if required by the target. results = target.search(query=rule.query_string, max_results=100) for text in results: @@ -51,4 +47,4 @@ def scan_task(scan_id): # Persist the completion time of the scan scan.finished_at = timezone.now() scan.save() - print(f'Saved scan results') + print('Saved scan results') diff --git a/chirps/scan/tests.py b/chirps/scan/tests.py index 7ce503c2..da3b49c4 100644 --- a/chirps/scan/tests.py +++ b/chirps/scan/tests.py @@ -1,3 +1 @@ -from django.test import TestCase - -# Create your tests here. +"""Tests for the scan application.""" diff --git a/chirps/scan/urls.py b/chirps/scan/urls.py index a04afd99..fc4067ff 100644 --- a/chirps/scan/urls.py +++ b/chirps/scan/urls.py @@ -1,3 +1,4 @@ +"""URLS for the scan application.""" from django.urls import path from . import views diff --git a/chirps/scan/views.py b/chirps/scan/views.py index bbd8a064..d970a48f 100644 --- a/chirps/scan/views.py +++ b/chirps/scan/views.py @@ -1,24 +1,28 @@ -import json +"""Views for the scan application.""" + from django.contrib.auth.decorators import login_required -from django.shortcuts import redirect, render, get_object_or_404 +from django.shortcuts import get_object_or_404, redirect, render from .forms import ScanForm -from .models import Result, Scan, Finding +from .models import Finding, Result, Scan from .tasks import scan_task + @login_required def finding_detail(request, finding_id): + """Render the finding detail page.""" finding = get_object_or_404(Finding, pk=finding_id, result__scan__user=request.user) return render(request, 'scan/finding_detail.html', {'finding': finding}) @login_required def result_detail(request, result_id): + """Render the scan result detail page.""" result = get_object_or_404(Result, pk=result_id, scan__user=request.user) return render(request, 'scan/result_detail.html', {'result': result}) @login_required def create(request): - + """Render the scan creation page and handle POST requests.""" if request.method == 'POST': scan_form = ScanForm(request.POST) if scan_form.is_valid(): @@ -50,7 +54,8 @@ def create(request): @login_required def dashboard(request): - # TODO: Add pagination + """Render the scan dashboard.""" + user_scans = Scan.objects.filter(user=request.user) # We're going to perform some manual aggregation (sqlite doesn't support calls to distinct()) @@ -60,7 +65,11 @@ def dashboard(request): for result in scan.result_set.all(): if result.rule.name not in scan.rules: - scan.rules[result.rule.name] = {'id': result.id, 'rule': result.rule, 'findings': Finding.objects.filter(result=result).count()} + scan.rules[result.rule.name] = { + 'id': result.id, + 'rule': result.rule, + 'findings': Finding.objects.filter(result=result).count() + } # Convert the dictionary into a list that the template can iterate on scan.rules = scan.rules.values() diff --git a/chirps/target/admin.py b/chirps/target/admin.py index 8c38f3f3..d611d5fe 100644 --- a/chirps/target/admin.py +++ b/chirps/target/admin.py @@ -1,3 +1,31 @@ +"""Admin interface definitions for target application models.""" + from django.contrib import admin +from polymorphic.admin import (PolymorphicChildModelAdmin, + PolymorphicParentModelAdmin) + +from .models import BaseTarget, MantiumTarget, PineconeTarget, RedisTarget + + +class BaseTargetAdmin(PolymorphicParentModelAdmin): + """Base admin class for the BaseTarget model.""" + + base_model = BaseTarget + + +class PineconeTargetAdmin(PolymorphicChildModelAdmin): + """Admin class for the PineconeTarget model.""" + base_model = PineconeTarget + +class MantiumTargetAdmin(PolymorphicChildModelAdmin): + """Admin class for the MantiumTarget model.""" + + base_model = MantiumTarget + +class RedisTargetAdmin(PolymorphicChildModelAdmin): + """Admin class for the RedisTarget model.""" + + base_model = RedisTarget -# Register your models here. +admin.site.register(RedisTarget) +admin.site.register(MantiumTarget) diff --git a/chirps/target/apps.py b/chirps/target/apps.py index 66dd4e3e..f5b7be2a 100644 --- a/chirps/target/apps.py +++ b/chirps/target/apps.py @@ -1,6 +1,9 @@ +"""Configuration options for the target application.""" from django.apps import AppConfig class TargetConfig(AppConfig): + """Configuration options for the target application.""" + default_auto_field = 'django.db.models.BigAutoField' name = 'target' diff --git a/chirps/target/custom_fields.py b/chirps/target/custom_fields.py index c05adf7b..2ed923b6 100644 --- a/chirps/target/custom_fields.py +++ b/chirps/target/custom_fields.py @@ -1,19 +1,23 @@ -# custom_fields.py - -from django.conf import settings -from fernet_fields import EncryptedCharField -from cryptography.fernet import Fernet - -class CustomEncryptedCharField(EncryptedCharField): - def __init__(self, *args, **kwargs): - self.fernet = Fernet(settings.FERNET_KEY) - super().__init__(*args, **kwargs) - - def from_db_value(self, value, expression, connection): - if value is not None: - value = super().from_db_value(value, expression, connection) - if isinstance(value, bytes): - return value.decode('utf-8') - else: - return value - return None +"""Custom fields for encrypting and decrypting data.""" + +from cryptography.fernet import Fernet +from django.conf import settings +from fernet_fields import EncryptedCharField + + +class CustomEncryptedCharField(EncryptedCharField): + """Custom encrypted char field that uses the fernet key from settings.""" + def __init__(self, *args, **kwargs): + """Initialize the field.""" + self.fernet = Fernet(settings.FERNET_KEY) + super().__init__(*args, **kwargs) + + def from_db_value(self, value, expression, connection, *args): + """Decrypt the value from the database.""" + if value is not None: + value = super().from_db_value(value, expression, connection, args) + if isinstance(value, bytes): + return value.decode('utf-8') + + return value + return None diff --git a/chirps/target/forms.py b/chirps/target/forms.py index 3890550b..03a328a3 100644 --- a/chirps/target/forms.py +++ b/chirps/target/forms.py @@ -1,11 +1,15 @@ +"""Forms for rendering and validating the target models.""" from django import forms from django.forms import ModelForm -from .models import MantiumTarget, RedisTarget, PineconeTarget +from .models import MantiumTarget, PineconeTarget, RedisTarget class RedisTargetForm(ModelForm): + """Form for the RedisTarget model.""" + class Meta: + """Django Meta options for the RedisTargetForm.""" model = RedisTarget fields = ['name', 'host', 'port', 'database_name', 'username', 'password'] @@ -20,7 +24,10 @@ class Meta: class MantiumTargetForm(ModelForm): + """Form for the MantiumTarget model.""" + class Meta: + """Django Meta options for the MantiumTargetForm.""" model = MantiumTarget fields = [ 'name', @@ -36,31 +43,33 @@ class Meta: 'client_secret': forms.PasswordInput(attrs={'class': 'form-control'}), } -class PineconeTargetForm(ModelForm): - class Meta: - model = PineconeTarget - fields = [ - 'name', - 'api_key', - 'environment', - 'index_name', - 'project_name', - ] - - widgets = { - 'name': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Enter a name for the target'}), - 'api_key': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'API Key'}), - 'environment': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Environment (optional)'}), - 'index_name': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Index Name (optional)'}), - 'project_name': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Project Name (optional)'}), - } - - -targets = [ - {'form': RedisTargetForm, 'model': RedisTarget}, - {'form': MantiumTargetForm, 'model': MantiumTarget}, - {'form': PineconeTargetForm, 'model': PineconeTarget}, -] +class PineconeTargetForm(ModelForm): + """Form for the PineconeTarget model.""" + class Meta: + """Django Meta options for the PineconeTargetForm.""" + model = PineconeTarget + fields = [ + 'name', + 'api_key', + 'environment', + 'index_name', + 'project_name', + ] + + widgets = { + 'name': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Enter a name for the target'}), + 'api_key': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'API Key'}), + 'environment': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Environment (optional)'}), + 'index_name': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Index Name (optional)'}), + 'project_name': forms.TextInput(attrs={'class': 'form-control', 'placeholder': 'Project Name (optional)'}), + } + + +targets = [ + {'form': RedisTargetForm, 'model': RedisTarget}, + {'form': MantiumTargetForm, 'model': MantiumTarget}, + {'form': PineconeTargetForm, 'model': PineconeTarget}, +] @@ -70,3 +79,4 @@ def target_from_html_name(html_name: str) -> dict: if target['model'].html_name == html_name: return target + return {} diff --git a/chirps/target/migrations/0001_initial.py b/chirps/target/migrations/0001_initial.py index d4bfa4e5..ed70ab89 100644 --- a/chirps/target/migrations/0001_initial.py +++ b/chirps/target/migrations/0001_initial.py @@ -1,9 +1,9 @@ # Generated by Django 4.2.2 on 2023-06-29 13:56 -from django.conf import settings -from django.db import migrations, models import django.db.models.deletion import fernet_fields.fields +from django.conf import settings +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/chirps/target/migrations/0002_pineconetarget.py b/chirps/target/migrations/0002_pineconetarget.py index aa8f85af..a5e7f40d 100644 --- a/chirps/target/migrations/0002_pineconetarget.py +++ b/chirps/target/migrations/0002_pineconetarget.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.3 on 2023-07-05 15:22 -from django.db import migrations, models import django.db.models.deletion +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/chirps/target/migrations/0003_alter_pineconetarget_api_key.py b/chirps/target/migrations/0003_alter_pineconetarget_api_key.py index fb6a2678..a1bc85dc 100644 --- a/chirps/target/migrations/0003_alter_pineconetarget_api_key.py +++ b/chirps/target/migrations/0003_alter_pineconetarget_api_key.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.3 on 2023-07-05 16:14 -from django.db import migrations import fernet_fields.fields +from django.db import migrations class Migration(migrations.Migration): diff --git a/chirps/target/migrations/0005_alter_pineconetarget_api_key.py b/chirps/target/migrations/0005_alter_pineconetarget_api_key.py index a798e37a..14f4bfa5 100644 --- a/chirps/target/migrations/0005_alter_pineconetarget_api_key.py +++ b/chirps/target/migrations/0005_alter_pineconetarget_api_key.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.3 on 2023-07-05 16:46 -from django.db import migrations import fernet_fields.fields +from django.db import migrations class Migration(migrations.Migration): diff --git a/chirps/target/migrations/0006_alter_pineconetarget_api_key.py b/chirps/target/migrations/0006_alter_pineconetarget_api_key.py index 928cadc8..ea73e544 100644 --- a/chirps/target/migrations/0006_alter_pineconetarget_api_key.py +++ b/chirps/target/migrations/0006_alter_pineconetarget_api_key.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.3 on 2023-07-05 16:51 -from django.db import migrations import target.models +from django.db import migrations class Migration(migrations.Migration): diff --git a/chirps/target/migrations/0007_alter_pineconetarget_api_key.py b/chirps/target/migrations/0007_alter_pineconetarget_api_key.py index 32f92e94..745b2b57 100644 --- a/chirps/target/migrations/0007_alter_pineconetarget_api_key.py +++ b/chirps/target/migrations/0007_alter_pineconetarget_api_key.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.3 on 2023-07-05 16:52 -from django.db import migrations import fernet_fields.fields +from django.db import migrations class Migration(migrations.Migration): diff --git a/chirps/target/migrations/0008_alter_pineconetarget_api_key.py b/chirps/target/migrations/0008_alter_pineconetarget_api_key.py index 440f124f..d183001a 100644 --- a/chirps/target/migrations/0008_alter_pineconetarget_api_key.py +++ b/chirps/target/migrations/0008_alter_pineconetarget_api_key.py @@ -1,7 +1,7 @@ # Generated by Django 4.2.3 on 2023-07-05 17:40 -from django.db import migrations import target.custom_fields +from django.db import migrations class Migration(migrations.Migration): diff --git a/chirps/target/models.py b/chirps/target/models.py index 7dad0fff..786457ad 100644 --- a/chirps/target/models.py +++ b/chirps/target/models.py @@ -1,23 +1,23 @@ """Models for the target appliation.""" -import pinecone -from django.contrib import admin +import pinecone +from django.contrib.auth.models import User from django.db import models +from django.templatetags.static import static from fernet_fields import EncryptedCharField from mantium_client.api_client import MantiumClient from mantium_spec.api.applications_api import ApplicationsApi -from polymorphic.admin import PolymorphicChildModelAdmin, PolymorphicParentModelAdmin from polymorphic.models import PolymorphicModel + from .custom_fields import CustomEncryptedCharField -from django.contrib.auth.models import User -from django.templatetags.static import static class BaseTarget(PolymorphicModel): """Base class that all targets will inherit from.""" name = models.CharField(max_length=128) user = models.ForeignKey(User, on_delete=models.CASCADE, null=True) + html_logo = None def search(self, query: str, max_results: int) -> list[str]: """Perform a query against the specified target, returning the max_results number of matches.""" @@ -26,15 +26,12 @@ def test_connection(self) -> bool: """Verify that the target can be connected to.""" def logo_url(self) -> str: + """Fetch the logo URL for the target.""" return static(self.html_logo) - def __str__(self): - return self.name - - -class BaseTargetAdmin(PolymorphicParentModelAdmin): - base_model = BaseTarget - + def __str__(self) -> str: + """String representation of this model.""" + return str(self.name) class RedisTarget(BaseTarget): """Implementation of a Redis target.""" @@ -61,58 +58,52 @@ def test_connection(self) -> bool: return True -class RedisTargetAdmin(PolymorphicChildModelAdmin): - base_model = RedisTarget - -class PineconeTarget(BaseTarget): - """Implementation of a Pinecone target.""" - - api_key = CustomEncryptedCharField(max_length=256, editable=True) - environment = models.CharField(max_length=256, blank=True, null=True) - index_name = models.CharField(max_length=256, blank=True, null=True) - project_name = models.CharField(max_length=256, blank=True, null=True) - - # Name of the file in the ./target/static/ directory to use as a logo - html_logo = 'target/pinecone-logo.png' - html_name = 'Pinecone' +class PineconeTarget(BaseTarget): + """Implementation of a Pinecone target.""" + + api_key = CustomEncryptedCharField(max_length=256, editable=True) + environment = models.CharField(max_length=256, blank=True, null=True) + index_name = models.CharField(max_length=256, blank=True, null=True) + project_name = models.CharField(max_length=256, blank=True, null=True) + + # Name of the file in the ./target/static/ directory to use as a logo + html_logo = 'target/pinecone-logo.png' + html_name = 'Pinecone' html_description = 'Pinecone Vector Database' - @property - def decrypted_api_key(self): - if self.api_key is not None: - try: - decrypted_value = self.api_key - return decrypted_value - except UnicodeDecodeError: - return "Error: Decryption failed" - return None - - def search(self, query: str, max_results: int) -> list[str]: - """Search the Pinecone target with the specified query.""" - pinecone.init(api_key=self.api_key, environment=self.environment) - - # Assuming the query is converted to a vector of the same dimension as the index. We should re-visit this. - query_vector = convert_query_to_vector(query) - - # Perform search on the Pinecone index - search_results = pinecone.fetch(index_name=self.index_name, query_vector=query_vector, top_k=max_results) - pinecone.deinit() - return search_results - - def test_connection(self) -> bool: - """Ensure that the Pinecone target can be connected to.""" - try: - pinecone.init(api_key=self.api_key, environment=self.environment) - - index_description = pinecone.describe_index(self.index_name) - pinecone.deinit() - return True - except Exception as e: - print(f"Pinecone connection test failed: {e}") + @property + def decrypted_api_key(self): + """Return the decrypted API key.""" + if self.api_key is not None: + try: + decrypted_value = self.api_key + return decrypted_value + except UnicodeDecodeError: + return "Error: Decryption failed" + return None + + def search(self, query: str, max_results: int) -> list[str]: + """Search the Pinecone target with the specified query.""" + pinecone.init(api_key=self.api_key, environment=self.environment) + + # Assuming the query is converted to a vector of the same dimension as the index. We should re-visit this. + query_vector = convert_query_to_vector(query) # pylint: disable=undefined-variable + + # Perform search on the Pinecone index + search_results = pinecone.fetch(index_name=self.index_name, query_vector=query_vector, top_k=max_results) + pinecone.deinit() + return search_results + + def test_connection(self) -> bool: + """Ensure that the Pinecone target can be connected to.""" + try: + pinecone.init(api_key=self.api_key, environment=self.environment) + pinecone.deinit() + return True + except Exception as err: # pylint: disable=broad-exception-caught + print(f"Pinecone connection test failed: {err}") return False -class PineconeTargetAdmin(PolymorphicChildModelAdmin): - base_model = PineconeTarget class MantiumTarget(BaseTarget): """Implementation of a Mantium target.""" @@ -137,12 +128,4 @@ def search(self, query: str, max_results: int) -> list[str]: documents = [doc['content'] for doc in results['documents']] return documents - -class MantiumTargetAdmin(PolymorphicChildModelAdmin): - base_model = MantiumTarget - -admin.site.register(RedisTarget) -admin.site.register(MantiumTarget) - -targets = [RedisTarget, MantiumTarget, PineconeTarget] - +targets = [RedisTarget, MantiumTarget, PineconeTarget] diff --git a/chirps/target/tests.py b/chirps/target/tests.py index b36523fe..5f032a89 100644 --- a/chirps/target/tests.py +++ b/chirps/target/tests.py @@ -1,11 +1,15 @@ +"""Test cases for the target application.""" +from django.contrib.auth.models import User # noqa: E5142 from django.test import TestCase from django.urls import reverse from target.models import MantiumTarget -from django.contrib.auth.models import User + class TargetTests(TestCase): + """Test the target application.""" def setUp(self): + """Initialize the database with some dummy users.""" self.users = [ {'username': 'user1', 'email': 'user1@mantiumai.com', 'password': 'user1password'}, {'username': 'user2', 'email': 'user2@mantiumai.com', 'password': 'user2password'}, @@ -25,7 +29,8 @@ def test_target_tenant_isolation(self): """Verify that targets are isolated to a single tenant.""" # Create a target for user1 - MantiumTarget.objects.create(name='Mantium Target', app_id='12345', token='1234', user=User.objects.get(username='user1')) + MantiumTarget.objects.create(name='Mantium Target', app_id='12345', client_id='1234', + client_secret='secret_dummy_value', user=User.objects.get(username='user1')) # Verify that the target is accessible to user1 (need to login first) response = self.client.post( diff --git a/chirps/target/urls.py b/chirps/target/urls.py index fc6f0411..3191ed03 100644 --- a/chirps/target/urls.py +++ b/chirps/target/urls.py @@ -1,3 +1,5 @@ +"""URLs for the target app.""" + from django.urls import path from . import views @@ -6,6 +8,5 @@ path('', views.dashboard, name='target_dashboard'), path('create/', views.create, name='target_create'), path('delete/', views.delete, name='target_delete'), - path('decrypted_keys/', views.decrypted_keys, name='decrypted_keys'), - + path('decrypted_keys/', views.decrypted_keys, name='decrypted_keys'), ] diff --git a/chirps/target/views.py b/chirps/target/views.py index ae492066..5918a335 100644 --- a/chirps/target/views.py +++ b/chirps/target/views.py @@ -1,28 +1,30 @@ """View handlers for targets.""" from django.contrib.auth.decorators import login_required +from django.http import JsonResponse from django.shortcuts import get_object_or_404, redirect, render -from django.http import JsonResponse from .forms import target_from_html_name, targets from .models import BaseTarget, PineconeTarget -def decrypted_keys(request): - keys = [] - for target in PineconeTarget.objects.all(): - keys.append(target.decrypted_api_key) - return JsonResponse({'keys': keys}) - -@login_required -def dashboard(request): - """Render the dashboard for the target app. - - Args: - request (HttpRequest): Django request object - """ - user_targets = BaseTarget.objects.filter(user=request.user) - return render( - request, 'target/dashboard.html', {'available_targets': targets, 'user_targets': user_targets} - ) + +def decrypted_keys(request): + """Return a list of decrypted API keys for all Pinecone targets.""" + keys = [] + for target in PineconeTarget.objects.all(): + keys.append(target.decrypted_api_key) + return JsonResponse({'keys': keys}) + +@login_required +def dashboard(request): + """Render the dashboard for the target app. + + Args: + request (HttpRequest): Django request object + """ + user_targets = BaseTarget.objects.filter(user=request.user) + return render( + request, 'target/dashboard.html', {'available_targets': targets, 'user_targets': user_targets} + ) @login_required @@ -63,7 +65,7 @@ def create(request, html_name): @login_required -def delete(request, target_id): +def delete(request, target_id): # pylint: disable=unused-argument """Delete a target from the database.""" get_object_or_404(BaseTarget, pk=target_id).delete() return redirect('target_dashboard')