From 293ad3a221462460aaa7c8f5fe4ef49cda10d02b Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Mon, 11 Sep 2023 14:47:05 +0800 Subject: [PATCH 1/3] feat: add project token --- dongtai_common/common/utils/__init__.py | 36 +++++++++++++++++++ .../migrations/0022_iastproject_token.py | 36 +++++++++++++++++++ dongtai_common/models/project.py | 3 ++ dongtai_common/models/user.py | 1 + dongtai_protocol/views/agent_register.py | 3 +- dongtai_web/urls.py | 2 ++ dongtai_web/views/project_token.py | 24 +++++++++++++ 7 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 dongtai_common/migrations/0022_iastproject_token.py create mode 100644 dongtai_web/views/project_token.py diff --git a/dongtai_common/common/utils/__init__.py b/dongtai_common/common/utils/__init__.py index 78703b80e..6b2561ab7 100644 --- a/dongtai_common/common/utils/__init__.py +++ b/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) diff --git a/dongtai_common/migrations/0022_iastproject_token.py b/dongtai_common/migrations/0022_iastproject_token.py new file mode 100644 index 000000000..66f86888b --- /dev/null +++ b/dongtai_common/migrations/0022_iastproject_token.py @@ -0,0 +1,36 @@ +# Generated by Django 3.2.20 on 2023-09-11 14:16 + +from django.db import migrations +import shortuuid.django_fields +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), + ] diff --git a/dongtai_common/models/project.py b/dongtai_common/models/project.py index b54e4469d..8efedfe26 100644 --- a/dongtai_common/models/project.py +++ b/dongtai_common/models/project.py @@ -9,6 +9,8 @@ from dongtai_common.models.strategy_user import IastStrategyUser from dongtai_common.utils.db import get_timestamp from dongtai_common.utils.settings import get_managed +from shortuuid.django_fields import ShortUUIDField +import string class VulValidation(models.IntegerChoices): @@ -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() diff --git a/dongtai_common/models/user.py b/dongtai_common/models/user.py index 30cf92f27..25428255b 100644 --- a/dongtai_common/models/user.py +++ b/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" diff --git a/dongtai_protocol/views/agent_register.py b/dongtai_protocol/views/agent_register.py index dcb1a2e71..94923c3a5 100644 --- a/dongtai_protocol/views/agent_register.py +++ b/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, diff --git a/dongtai_web/urls.py b/dongtai_web/urls.py index 9799f2221..997e4e15f 100644 --- a/dongtai_web/urls.py +++ b/dongtai_web/urls.py @@ -125,6 +125,7 @@ from dongtai_web.vul_log.vul_log_view import VulLogViewSet from dongtai_web.vul_recheck_payload.vul_recheck_payload import VulReCheckPayloadViewSet from static.i18n.views.setlang import LanguageSetting +from dongtai_web.views.project_token import ProjectToken urlpatterns: list[URLResolver | URLPattern] = [ path("user/", UserDetailEndPoint.as_view()), @@ -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//token", ProjectToken.as_view()), path("project/", ProjectDetail.as_view()), path("project/add", ProjectAdd.as_view()), path("project/delete", ProjectDel.as_view()), diff --git a/dongtai_web/views/project_token.py b/dongtai_web/views/project_token.py new file mode 100644 index 000000000..04a687f28 --- /dev/null +++ b/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")) From 5a62e328a8c957bff9d1a1847dd5d3de2048b228 Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Mon, 11 Sep 2023 14:49:01 +0800 Subject: [PATCH 2/3] feat: add project token --- dongtai_conf/settings.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dongtai_conf/settings.py b/dongtai_conf/settings.py index 28159e59f..ff8aa7f8a 100644 --- a/dongtai_conf/settings.py +++ b/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": [ From 5765508987aab5b23e6a5b29ac2cfc35575d2196 Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Mon, 11 Sep 2023 14:50:58 +0800 Subject: [PATCH 3/3] feat: add project token --- dongtai_common/migrations/0022_iastproject_token.py | 2 +- dongtai_common/models/project.py | 4 ++-- dongtai_web/urls.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dongtai_common/migrations/0022_iastproject_token.py b/dongtai_common/migrations/0022_iastproject_token.py index 66f86888b..20083d2e7 100644 --- a/dongtai_common/migrations/0022_iastproject_token.py +++ b/dongtai_common/migrations/0022_iastproject_token.py @@ -1,7 +1,7 @@ # Generated by Django 3.2.20 on 2023-09-11 14:16 -from django.db import migrations import shortuuid.django_fields +from django.db import migrations from shortuuid import ShortUUID diff --git a/dongtai_common/models/project.py b/dongtai_common/models/project.py index 8efedfe26..ddfbc87a8 100644 --- a/dongtai_common/models/project.py +++ b/dongtai_common/models/project.py @@ -1,16 +1,16 @@ #!/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 from dongtai_common.models.strategy_user import IastStrategyUser from dongtai_common.utils.db import get_timestamp from dongtai_common.utils.settings import get_managed -from shortuuid.django_fields import ShortUUIDField -import string class VulValidation(models.IntegerChoices): diff --git a/dongtai_web/urls.py b/dongtai_web/urls.py index 997e4e15f..2359b7259 100644 --- a/dongtai_web/urls.py +++ b/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 @@ -125,7 +126,6 @@ from dongtai_web.vul_log.vul_log_view import VulLogViewSet from dongtai_web.vul_recheck_payload.vul_recheck_payload import VulReCheckPayloadViewSet from static.i18n.views.setlang import LanguageSetting -from dongtai_web.views.project_token import ProjectToken urlpatterns: list[URLResolver | URLPattern] = [ path("user/", UserDetailEndPoint.as_view()),