Merge pull request #1371 from pierotofy/quotas

External auth support, task sizes, quotas

.env

@@ -10,3 +10,4 @@ WO_DEBUG=NO
 WO_DEV=NO
 WO_BROKER=redis://broker
 WO_DEFAULT_NODES=1
+WO_SETTINGS=

app/admin.py

@@ -10,10 +10,13 @@
 from django.urls import reverse
 from django.utils.html import format_html
 from guardian.admin import GuardedModelAdmin
+from django.contrib.auth.admin import UserAdmin as BaseUserAdmin
+from django.contrib.auth.models import User
 
 from app.models import PluginDatum
 from app.models import Preset
 from app.models import Plugin
+from app.models import Profile
 from app.plugins import get_plugin_by_name, enable_plugin, disable_plugin, delete_plugin, valid_plugin, \
     get_plugins_persistent_path, clear_plugins_cache, init_plugins
 from .models import Project, Task, Setting, Theme
@@ -260,3 +263,14 @@ def plugin_actions(self, obj):
 
 
 admin.site.register(Plugin, PluginAdmin)
+
+class ProfileInline(admin.StackedInline):
+    model = Profile
+    can_delete = False
+
+class UserAdmin(BaseUserAdmin):
+    inlines = [ProfileInline]
+
+# Re-register UserAdmin
+admin.site.unregister(User)
+admin.site.register(User, UserAdmin)

app/api/admin.py

@@ -1,7 +1,10 @@
 from django.contrib.auth.models import User, Group
-from rest_framework import serializers, viewsets, generics, status
+from app.models import Profile
+from rest_framework import serializers, viewsets, generics, status, exceptions
+from rest_framework.decorators import action
 from rest_framework.permissions import IsAdminUser
 from rest_framework.response import Response
+from django.core.exceptions import ObjectDoesNotExist
 from django.contrib.auth.hashers import make_password
 from app import models
 
@@ -20,6 +23,7 @@ def get_queryset(self):
         if email is not None:
             queryset = queryset.filter(email=email)
         return queryset
+
     def create(self, request):
         data = request.data.copy()
         password = data.get('password')
@@ -44,3 +48,37 @@ def get_queryset(self):
         if name is not None:
             queryset = queryset.filter(name=name)
         return queryset
+
+
+class ProfileSerializer(serializers.ModelSerializer):
+    class Meta:
+        model  = Profile
+        exclude = ('id', ) 
+
+        read_only_fields = ('user', )
+
+class AdminProfileViewSet(viewsets.ModelViewSet):
+    pagination_class = None
+    serializer_class = ProfileSerializer
+    permission_classes = [IsAdminUser]
+    lookup_field = 'user'
+
+    def get_queryset(self):
+        return Profile.objects.all()
+
+
+    @action(detail=True, methods=['post'])
+    def update_quota_deadline(self, request, user=None):
+        try:
+            hours = float(request.data.get('hours', ''))
+            if hours < 0:
+                raise ValueError("hours must be >= 0")
+        except ValueError as e:
+            raise exceptions.ValidationError(str(e))
+
+        try:
+            p = Profile.objects.get(user=user)
+        except ObjectDoesNotExist:
+            raise exceptions.NotFound()
+
+        return Response({'deadline': p.set_quota_deadline(hours)}, status=status.HTTP_200_OK)

app/api/externalauth.py

@@ -0,0 +1,39 @@
+from django.contrib.auth.models import User
+from django.contrib.auth import login
+from rest_framework.views import APIView
+from rest_framework import exceptions, permissions, parsers
+from rest_framework.response import Response
+from app.auth.backends import get_user_from_external_auth_response
+import requests
+from webodm import settings
+
+class ExternalTokenAuth(APIView):
+    permission_classes = (permissions.AllowAny,)
+    parser_classes = (parsers.JSONParser, parsers.FormParser,)
+
+    def post(self, request):
+        # This should never happen
+        if settings.EXTERNAL_AUTH_ENDPOINT == '':
+            return Response({'error': 'EXTERNAL_AUTH_ENDPOINT not set'})
+
+        token = request.COOKIES.get('external_access_token', '')
+        if token == '':
+            return Response({'error': 'external_access_token cookie not set'})
+
+        try:
+            r = requests.post(settings.EXTERNAL_AUTH_ENDPOINT, headers={
+                'Authorization': "Bearer %s" % token 
+            })
+            res = r.json()
+            if res.get('user_id') is not None:
+                user = get_user_from_external_auth_response(res)
+                if user is not None:
+                    login(request, user, backend='django.contrib.auth.backends.ModelBackend')
+                    return Response({'redirect': '/'})
+                else:
+                    return Response({'error': 'Invalid credentials'})
+            else:
+                return Response({'error': res.get('message', 'Invalid external server response')})
+        except Exception as e:
+            return Response({'error': str(e)})
+

app/api/tasks.py

@@ -75,7 +75,7 @@ def get_can_rerun_from(self, obj):
     class Meta:
         model = models.Task
         exclude = ('console_output', 'orthophoto_extent', 'dsm_extent', 'dtm_extent', )
-        read_only_fields = ('processing_time', 'status', 'last_error', 'created_at', 'pending_action', 'available_assets', )
+        read_only_fields = ('processing_time', 'status', 'last_error', 'created_at', 'pending_action', 'available_assets', 'size', )
 
 class TaskViewSet(viewsets.ViewSet):
     """
@@ -184,6 +184,7 @@ def commit(self, request, pk=None, project_pk=None):
         if task.images_count < 1:
             raise exceptions.ValidationError(detail=_("You need to upload at least 1 file before commit"))
 
+        task.update_size()
         task.save()
         worker_tasks.process_task.delay(task.id)
 

app/api/urls.py

@@ -6,13 +6,14 @@
 from .tasks import TaskViewSet, TaskDownloads, TaskAssets, TaskAssetsImport
 from .imageuploads import Thumbnail, ImageDownload
 from .processingnodes import ProcessingNodeViewSet, ProcessingNodeOptionsView
-from .admin import AdminUserViewSet, AdminGroupViewSet
+from .admin import AdminUserViewSet, AdminGroupViewSet, AdminProfileViewSet
 from rest_framework_nested import routers
 from rest_framework_jwt.views import obtain_jwt_token
 from .tiler import TileJson, Bounds, Metadata, Tiles, Export
 from .potree import Scene, CameraView
 from .workers import CheckTask, GetTaskResult
 from .users import UsersList
+from .externalauth import ExternalTokenAuth
 from webodm import settings
 
 router = routers.DefaultRouter()
@@ -26,6 +27,7 @@
 admin_router = routers.DefaultRouter()
 admin_router.register(r'admin/users', AdminUserViewSet, basename='admin-users')
 admin_router.register(r'admin/groups', AdminGroupViewSet, basename='admin-groups')
+admin_router.register(r'admin/profiles', AdminProfileViewSet, basename='admin-groups')
 
 urlpatterns = [
     url(r'processingnodes/options/$', ProcessingNodeOptionsView.as_view()),
@@ -56,9 +58,12 @@
     url(r'^auth/', include('rest_framework.urls')),
     url(r'^token-auth/', obtain_jwt_token),
 
-    url(r'^plugins/(?P<plugin_name>[^/.]+)/(.*)$', api_view_handler)
+    url(r'^plugins/(?P<plugin_name>[^/.]+)/(.*)$', api_view_handler),
 ]
 
 if settings.ENABLE_USERS_API:
     urlpatterns.append(url(r'users', UsersList.as_view()))
 
+if settings.EXTERNAL_AUTH_ENDPOINT != '':
+    urlpatterns.append(url(r'^external-token-auth/', ExternalTokenAuth.as_view()))
+

app/auth/backends.py

@@ -0,0 +1,88 @@
+import requests
+from django.contrib.auth.backends import ModelBackend
+from django.contrib.auth.models import User
+from nodeodm.models import ProcessingNode
+from webodm import settings
+from guardian.shortcuts import assign_perm
+import logging
+
+logger = logging.getLogger('app.logger')
+
+def get_user_from_external_auth_response(res):
+    if 'message' in res or 'error' in res:
+        return None
+
+    if 'user_id' in res and 'username' in res:
+        try:
+            user = User.objects.get(pk=res['user_id'])
+        except User.DoesNotExist:
+            user = User(pk=res['user_id'], username=res['username'])
+            user.save()
+
+        # Update user info
+        if user.username != res['username']:
+            user.username = res['username']
+            user.save()
+
+        maxQuota = -1
+        if 'maxQuota' in res:
+            maxQuota = res['maxQuota']
+        if 'node' in res and 'limits' in res['node'] and 'maxQuota' in res['node']['limits']:
+            maxQuota = res['node']['limits']['maxQuota']
+
+        # Update quotas
+        if user.profile.quota != maxQuota:
+            user.profile.quota = maxQuota
+            user.save()
+
+        # Setup/update processing node
+        if 'node' in res and 'hostname' in res['node'] and 'port' in res['node']:
+            hostname = res['node']['hostname']
+            port = res['node']['port']
+            token = res['node'].get('token', '')
+
+            # Only add/update if a token is provided, since we use 
+            # tokens as unique identifiers for hostname/port updates
+            if token != "":
+                try:
+                    node = ProcessingNode.objects.get(token=token)
+                    if node.hostname != hostname or node.port != port:
+                        node.hostname = hostname
+                        node.port = port
+                        node.save()
+
+                except ProcessingNode.DoesNotExist:
+                    node = ProcessingNode(hostname=hostname, port=port, token=token)
+                    node.save()
+
+                if not user.has_perm('view_processingnode', node):
+                    assign_perm('view_processingnode', user, node)
+
+        return user
+    else:
+        return None
+
+class ExternalBackend(ModelBackend):
+    def authenticate(self, request, username=None, password=None):
+        if settings.EXTERNAL_AUTH_ENDPOINT == "":
+            return None
+
+        try:
+            r = requests.post(settings.EXTERNAL_AUTH_ENDPOINT, {
+                'username': username,
+                'password': password
+            }, headers={'Accept': 'application/json'})
+            res = r.json()
+
+            return get_user_from_external_auth_response(res)
+        except:
+            return None
+
+    def get_user(self, user_id):
+        if settings.EXTERNAL_AUTH_ENDPOINT == "":
+            return None
+
+        try:
+            return User.objects.get(pk=user_id)
+        except User.DoesNotExist:
+            return None

app/migrations/0036_task_size.py

@@ -0,0 +1,50 @@
+# Generated by Django 2.2.27 on 2023-08-21 14:50
+import os
+from django.db import migrations, models
+from webodm import settings
+
+def task_path(project_id, task_id, *args):
+    return os.path.join(settings.MEDIA_ROOT,
+                        "project",
+                        str(project_id),
+                        "task",
+                        str(task_id),
+                        *args)
+
+def update_size(task):
+    try:
+        total_bytes = 0
+        for dirpath, _, filenames in os.walk(task_path(task.project.id, task.id)):
+            for f in filenames:
+                fp = os.path.join(dirpath, f)
+                if not os.path.islink(fp):
+                    total_bytes += os.path.getsize(fp)
+        task.size = (total_bytes / 1024 / 1024)
+        task.save()
+        print("Updated {} with size {}".format(task, task.size))
+    except Exception as e:
+        print("Cannot update size for task {}: {}".format(task, str(e)))
+
+
+
+def update_task_sizes(apps, schema_editor):
+    Task = apps.get_model('app', 'Task')
+
+    for t in Task.objects.all():
+        update_size(t)
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('app', '0035_task_orthophoto_bands'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='task',
+            name='size',
+            field=models.FloatField(blank=True, default=0.0, help_text='Size of the task on disk in megabytes', verbose_name='Size'),
+        ),
+
+        migrations.RunPython(update_task_sizes),
+    ]

app/migrations/0037_profile.py

@@ -0,0 +1,35 @@
+# Generated by Django 2.2.27 on 2023-08-24 16:35
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+def create_profiles(apps, schema_editor):
+    User = apps.get_model('auth', 'User')
+    Profile = apps.get_model('app', 'Profile')
+
+    for u in User.objects.all():
+        p = Profile.objects.create(user=u)
+        p.save()
+        print("Created user profile for %s" % u.username)
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('app', '0036_task_size'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='Profile',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('quota', models.FloatField(blank=True, default=-1, help_text='Maximum disk quota in megabytes', verbose_name='Quota')),
+                ('user', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)),
+            ],
+        ),
+
+        migrations.RunPython(create_profiles),
+    ]

app/models/__init__.py

@@ -5,6 +5,7 @@
 from .setting import Setting
 from .plugin_datum import PluginDatum
 from .plugin import Plugin
+from .profile import Profile
 
 # deprecated
 def image_directory_path(image_upload, filename):