Merge pull request #1797 from Bidaya0/feat/add-project-token

Feat/add project token

dongtai_common/common/utils/__init__.py

@@ -159,3 +159,39 @@ def authenticate(self, request):
             return None
         token = auth.lower().replace(self.keyword.lower().encode(), b"", 1).decode()
         return self.auth_decodedenticate_credentials(token)
+
+
+@cached_decorator(random_range=(60, 120), use_celery_update=False)
+def get_user_from_project_key(key):
+    from dongtai_common.models.project import IastProject
+    from dongtai_common.models.user import User
+
+    project = IastProject.objects.get(token=key)
+    principal = User.objects.filter(pk=project.user_id).first()
+    user = principal if principal else User.objects.filter(pk=1).first()
+    user.using_project = project
+    return user
+
+
+class ProjectTokenAuthentication(TokenAuthentication):
+    keyword = "Token PROJECT"
+    model = None
+
+    def auth_decodedenticate_credentials(self, key):
+        from rest_framework import exceptions
+
+        from dongtai_common.models.project import IastProject
+
+        try:
+            warnings.warn("Department token is departured. Please use new token to register agent.", stacklevel=1)
+            user = get_user_from_project_key(key)
+        except IastProject.DoesNotExist as e:
+            raise exceptions.AuthenticationFailed(_("Invalid token.")) from e
+        return (user, key)
+
+    def authenticate(self, request):
+        auth = get_authorization_header(request)
+        if not auth or not auth.lower().startswith(self.keyword.lower().encode()):
+            return None
+        token = auth.lower().replace(self.keyword.lower().encode(), b"", 1).decode()
+        return self.auth_decodedenticate_credentials(token)

dongtai_common/migrations/0022_iastproject_token.py

@@ -0,0 +1,36 @@
+# Generated by Django 3.2.20 on 2023-09-11 14:16
+
+import shortuuid.django_fields
+from django.db import migrations
+from shortuuid import ShortUUID
+
+
+def update_exist_project_token(apps, schema_editor):
+    IastProject = apps.get_model("dongtai_common", "IastProject")
+    objs_list = []
+    for project in IastProject.objects.all():
+        project.token = ShortUUID(alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789").random(
+            length=22
+        )
+        objs_list.append(project)
+    IastProject.objects.bulk_update(objs_list, ["token"])
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("dongtai_common", "0021_iastwebhooklog"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="iastproject",
+            name="token",
+            field=shortuuid.django_fields.ShortUUIDField(
+                alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789",
+                length=22,
+                max_length=22,
+                prefix="",
+            ),
+        ),
+        migrations.RunPython(update_exist_project_token),
+    ]

dongtai_common/models/project.py

@@ -1,8 +1,10 @@
 #!/usr/bin/env python
 # datetime:2020/11/30 下午5:32
+import string
 import time
 
 from django.db import models
+from shortuuid.django_fields import ShortUUIDField
 
 from dongtai_common.models import User
 from dongtai_common.models.department import Department
@@ -79,6 +81,7 @@ class IastProject(models.Model):
     status = models.IntegerField(default=0, choices=ProjectStatus.choices)
     projectgroups = models.ManyToManyField("IastProjectGroup", through="IastProjectGroupProject")
     users = models.ManyToManyField("User", through="IastProjectUser", related_name="auth_projects")
+    token = ShortUUIDField(max_length=22, alphabet=string.ascii_letters + string.digits)
 
     class Meta:
         managed = get_managed()

dongtai_common/models/user.py

@@ -71,6 +71,7 @@ class User(AbstractUser, PermissionsMixin):
 
     objects = SaaSUserManager()
     using_department = None
+    using_project = None
 
     class Meta(AbstractUser.Meta):
         db_table = "auth_user"

dongtai_conf/settings.py

@@ -121,6 +121,7 @@ def get_installed_apps():
     "DEFAULT_AUTHENTICATION_CLASSES": [
         "rest_framework.authentication.SessionAuthentication",
         "dongtai_common.common.utils.DepartmentTokenAuthentication",
+        "dongtai_common.common.utils.ProjectTokenAuthentication",
         "rest_framework.authentication.TokenAuthentication",
     ],
     "DEFAULT_RENDERER_CLASSES": [

dongtai_protocol/views/agent_register.py

@@ -265,7 +265,8 @@ def post(self, request: Request):
             }
 
             default_params.update(template.to_full_project_args() if template else {})
-
+            if request.user.using_project is not None:
+                project_name = request.user.using_project.name
             with transaction.atomic():
                 (
                     obj,

dongtai_web/urls.py

@@ -74,6 +74,7 @@
 from dongtai_web.views.project_engines import ProjectEngines
 from dongtai_web.views.project_search import ProjectSearch
 from dongtai_web.views.project_summary import ProjectSummary
+from dongtai_web.views.project_token import ProjectToken
 from dongtai_web.views.project_version_add import ProjectVersionAdd
 from dongtai_web.views.project_version_current import ProjectVersionCurrent
 from dongtai_web.views.project_version_delete import ProjectVersionDelete
@@ -137,6 +138,7 @@
     path("user/password/reset", UserPasswordReset.as_view()),
     path("captcha/", include("captcha.urls")),
     path(r"captcha/refresh", CaptchaCreate.as_view()),
+    path("project/<int:pk>/token", ProjectToken.as_view()),
     path("project/<int:id>", ProjectDetail.as_view()),
     path("project/add", ProjectAdd.as_view()),
     path("project/delete", ProjectDel.as_view()),

dongtai_web/views/project_token.py

@@ -0,0 +1,24 @@
+from django.utils.translation import gettext_lazy as _
+
+from dongtai_common.endpoint import R, UserEndPoint
+from dongtai_web.utils import extend_schema_with_envcheck
+
+
+class ProjectToken(UserEndPoint):
+    @extend_schema_with_envcheck(
+        tags=[_("Project")],
+        summary=_("Projects Token"),
+        description=_(
+            "Get project information by project id, including the current version information of the project."
+        ),
+    )
+    def get(self, request, pk):
+        project = request.user.get_projects().filter(pk=pk).first()
+        if project:
+            return R.success(
+                data={
+                    "id": project.id,
+                    "token": project.token,
+                }
+            )
+        return R.failure(status=203, msg=_("no permission"))