From a1ae29d35ebcce0d936a74e0b4d028c3b3af166f Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Wed, 14 Aug 2024 16:15:07 +0800
Subject: [PATCH] fix: Use only_sudo failed (#13966)

* fix: Use only_sudo failed

* fix: Use only_sudo failed

* fix: Use only_sudo failed

---------

Co-authored-by: feng <1304903146@qq.com>
---
 apps/accounts/models/account.py     |  3 ++-
 apps/assets/const/__init__.py       |  1 +
 apps/assets/const/platform.py       | 11 +++++++++++
 apps/assets/models/platform.py      | 13 ++++++++++++-
 apps/assets/serializers/platform.py | 15 +++------------
 5 files changed, 29 insertions(+), 14 deletions(-)
 create mode 100644 apps/assets/const/platform.py

diff --git a/apps/accounts/models/account.py b/apps/accounts/models/account.py
index 4a175609f..f99230670 100644
--- a/apps/accounts/models/account.py
+++ b/apps/accounts/models/account.py
@@ -119,7 +119,8 @@ class Account(AbsConnectivity, LabeledMixin, BaseAccount):
             return auth
 
         auth.update(self.make_account_ansible_vars(su_from))
-        become_method = platform.su_method if platform.su_method else 'sudo'
+
+        become_method = platform.ansible_become_method
         password = su_from.secret if become_method == 'sudo' else self.secret
         auth['ansible_become'] = True
         auth['ansible_become_method'] = become_method
diff --git a/apps/assets/const/__init__.py b/apps/assets/const/__init__.py
index 9e3f2cbb1..abf7e71b1 100644
--- a/apps/assets/const/__init__.py
+++ b/apps/assets/const/__init__.py
@@ -2,5 +2,6 @@ from .automation import *
 from .base import *
 from .category import *
 from .host import *
+from .platform import *
 from .protocol import *
 from .types import *
diff --git a/apps/assets/const/platform.py b/apps/assets/const/platform.py
new file mode 100644
index 000000000..f37d9fcc6
--- /dev/null
+++ b/apps/assets/const/platform.py
@@ -0,0 +1,11 @@
+from django.db.models import TextChoices
+
+
+class SuMethodChoices(TextChoices):
+    sudo = "sudo", "sudo su -"
+    su = "su", "su - "
+    only_sudo = "only_sudo", "sudo su"
+    only_su = "only_su", "su"
+    enable = "enable", "enable"
+    super = "super", "super 15"
+    super_level = "super_level", "super level 15"
diff --git a/apps/assets/models/platform.py b/apps/assets/models/platform.py
index 85ec8d535..d66f6bd15 100644
--- a/apps/assets/models/platform.py
+++ b/apps/assets/models/platform.py
@@ -1,7 +1,7 @@
 from django.db import models
 from django.utils.translation import gettext_lazy as _
 
-from assets.const import AllTypes, Category, Protocol
+from assets.const import AllTypes, Category, Protocol, SuMethodChoices
 from common.db.fields import JsonDictTextField
 from common.db.models import JMSBaseModel
 
@@ -127,6 +127,17 @@ class Platform(LabeledMixin, JMSBaseModel):
             return True
         return False
 
+    @property
+    def ansible_become_method(self):
+        su_method = self.su_method or SuMethodChoices.sudo
+        if su_method in [SuMethodChoices.sudo, SuMethodChoices.only_sudo]:
+            method = SuMethodChoices.sudo
+        elif su_method in [SuMethodChoices.su, SuMethodChoices.only_su]:
+            method = SuMethodChoices.su
+        else:
+            method = su_method
+        return method
+
     def __str__(self):
         return self.name
 
diff --git a/apps/assets/serializers/platform.py b/apps/assets/serializers/platform.py
index 93c0e54dd..0b538ee52 100644
--- a/apps/assets/serializers/platform.py
+++ b/apps/assets/serializers/platform.py
@@ -9,7 +9,7 @@ from common.serializers import (
 )
 from common.serializers.fields import LabeledChoiceField
 from common.utils import lazyproperty
-from ..const import Category, AllTypes, Protocol
+from ..const import Category, AllTypes, Protocol, SuMethodChoices
 from ..models import Platform, PlatformProtocol, PlatformAutomation
 
 __all__ = ["PlatformSerializer", "PlatformOpsMethodSerializer", "PlatformProtocolSerializer"]
@@ -124,15 +124,6 @@ class PlatformCustomField(serializers.Serializer):
 
 
 class PlatformSerializer(ResourceLabelsMixin, WritableNestedModelSerializer):
-    SU_METHOD_CHOICES = [
-        ("sudo", "sudo su -"),
-        ("su", "su - "),
-        ("only_sudo", "sudo su"),
-        ("only_su", "su"),
-        ("enable", "enable"),
-        ("super", "super 15"),
-        ("super_level", "super level 15")
-    ]
     id = serializers.IntegerField(
         label='ID', required=False,
         validators=[UniqueValidator(queryset=Platform.objects.all())]
@@ -143,8 +134,8 @@ class PlatformSerializer(ResourceLabelsMixin, WritableNestedModelSerializer):
     protocols = PlatformProtocolSerializer(label=_("Protocols"), many=True, required=False)
     automation = PlatformAutomationSerializer(label=_("Automation"), required=False, default=dict)
     su_method = LabeledChoiceField(
-        choices=SU_METHOD_CHOICES, label=_("Su method"),
-        required=False, default="sudo", allow_null=True
+        choices=SuMethodChoices.choices, label=_("Su method"),
+        required=False, default=SuMethodChoices.sudo, allow_null=True
     )
     custom_fields = PlatformCustomField(label=_("Custom fields"), many=True, required=False)