netcmd: move get_policy method from base class to the model

There isn't much left of the base class, the next thing is to remove
it.

Signed-off-by: Rob van der Linde <rob@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>
Reviewed-by: Joseph Sutton <josephsutton@catalyst.net.nz>

3 files changed, 71 insertions, 23 deletions

diff --git a/python/samba/netcmd/domain/auth/base.py b/python/samba/netcmd/domain/auth/base.py
index 1a3633d9f3b..a33e0703d3e 100644
--- a/python/samba/netcmd/domain/auth/base.py
+++ b/python/samba/netcmd/domain/auth/base.py
@@ -20,21 +20,10 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-from samba.netcmd import Command, CommandError
-from samba.netcmd.domain.models import AuthenticationPolicy
+from samba.netcmd import Command
 
 
 class SiloCommand(Command):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.ldb = None
-
-    def get_policy(self, name):
-        """Helper function to return auth policy or raise CommandError.
-
-        :raises CommandError: if the policy was not found.
-        """
-        policy = AuthenticationPolicy.get(self.ldb, cn=name)
-        if policy is None:
-            raise CommandError(f"Authentication policy {name} not found.")
-        return policy
diff --git a/python/samba/netcmd/domain/auth/silo.py b/python/samba/netcmd/domain/auth/silo.py
index a7017e1dfe1..df90180061f 100644
--- a/python/samba/netcmd/domain/auth/silo.py
+++ b/python/samba/netcmd/domain/auth/silo.py
@@ -23,7 +23,7 @@
 import samba.getopt as options
 from ldb import LdbError
 from samba.netcmd import CommandError, Option, SuperCommand
-from samba.netcmd.domain.models import AuthenticationSilo
+from samba.netcmd.domain.models import AuthenticationPolicy, AuthenticationSilo
 
 from .base import SiloCommand
 from .silo_member import cmd_domain_auth_silo_member
@@ -141,6 +141,18 @@ class cmd_domain_auth_silo_create(SiloCommand):
                dest="enforce", action="store_true")
     ]
 
+    @staticmethod
+    def get_policy(ldb, name):
+        """Helper function to fetch auth policy or raise CommandError.
+
+        :param ldb: Ldb connection
+        :param name: Either the DN or name of authentication policy
+        """
+        try:
+            return AuthenticationPolicy.lookup(ldb, name)
+        except (LookupError, ValueError) as e:
+            raise CommandError(e)
+
     def run(self, ldap_url=None, sambaopts=None, credopts=None, name=None,
             description=None, policy=None, user_policy=None,
             service_policy=None, computer_policy=None, protect=None,
@@ -172,15 +184,15 @@ class cmd_domain_auth_silo_create(SiloCommand):
 
         # Set user policy
         if user_policy:
-            silo.user_policy = self.get_policy(user_policy).dn
+            silo.user_policy = self.get_policy(self.ldb, user_policy).dn
 
         # Set service policy
         if service_policy:
-            silo.service_policy = self.get_policy(service_policy).dn
+            silo.service_policy = self.get_policy(self.ldb, service_policy).dn
 
         # Set computer policy
         if computer_policy:
-            silo.computer_policy = self.get_policy(computer_policy).dn
+            silo.computer_policy = self.get_policy(self.ldb, computer_policy).dn
 
         # Either --enforce will be set or --audit but never both.
         # The default if both are missing is enforce=True.
@@ -246,6 +258,18 @@ class cmd_domain_auth_silo_modify(SiloCommand):
                dest="enforce", action="store_true")
     ]
 
+    @staticmethod
+    def get_policy(ldb, name):
+        """Helper function to fetch auth policy or raise CommandError.
+
+        :param ldb: Ldb connection
+        :param name: Either the DN or name of authentication policy
+        """
+        try:
+            return AuthenticationPolicy.lookup(ldb, name)
+        except (LookupError, ValueError) as e:
+            raise CommandError(e)
+
     def run(self, ldap_url=None, sambaopts=None, credopts=None, name=None,
             description=None, policy=None, user_policy=None,
             service_policy=None, computer_policy=None, protect=None,
@@ -282,13 +306,23 @@ class cmd_domain_auth_silo_modify(SiloCommand):
         if description is not None:
             silo.description = description
 
-        # Silo policies.
-        if user_policy is not None:
-            silo.user_policy = self.get_policy(user_policy).dn
-        if service_policy is not None:
-            silo.service_policy = self.get_policy(service_policy).dn
-        if computer_policy is not None:
-            silo.computer_policy = self.get_policy(computer_policy).dn
+        # Set or unset user policy.
+        if user_policy == "":
+            silo.user_policy = None
+        elif user_policy:
+            silo.user_policy = self.get_policy(self.ldb, user_policy).dn
+
+        # Set or unset service policy.
+        if service_policy == "":
+            silo.service_policy = None
+        elif service_policy:
+            silo.service_policy = self.get_policy(self.ldb, service_policy).dn
+
+        # Set or unset computer policy.
+        if computer_policy == "":
+            silo.computer_policy = None
+        elif computer_policy:
+            silo.computer_policy = self.get_policy(self.ldb, computer_policy).dn
 
         # Update silo
         try:
diff --git a/python/samba/netcmd/domain/models/auth_policy.py b/python/samba/netcmd/domain/models/auth_policy.py
index fa0b07be910..dec8bb26190 100644
--- a/python/samba/netcmd/domain/models/auth_policy.py
+++ b/python/samba/netcmd/domain/models/auth_policy.py
@@ -21,6 +21,7 @@
 #
 
 from enum import IntEnum
+from ldb import Dn
 
 from .fields import BooleanField, EnumField, IntegerField, StringField
 from .model import Model
@@ -71,3 +72,27 @@ class AuthenticationPolicy(Model):
     @staticmethod
     def get_object_class():
         return "msDS-AuthNPolicy"
+
+    @staticmethod
+    def lookup(ldb, name):
+        """Helper function to return auth policy or raise LookupError.
+
+        :param ldb: Ldb connection
+        :param name: Either DN or name of Authentication Policy
+        :raises: LookupError if not found
+        :raises: ValueError if name is not set
+        """
+        if not name:
+            raise ValueError("Attribute 'name' is required.")
+
+        try:
+            # It's possible name is already a Dn.
+            dn = name if isinstance(name, Dn) else Dn(ldb, name)
+            policy = AuthenticationPolicy.get(ldb, dn=dn)
+        except ValueError:
+            policy = AuthenticationPolicy.get(ldb, cn=name)
+
+        if policy is None:
+            raise LookupError(f"Authentication policy {name} not found.")
+
+        return policy