diff --git a/enable_shell_access.py b/enable_shell_access.py
index 65e15bf..68e35c5 100644
--- a/enable_shell_access.py
+++ b/enable_shell_access.py
@@ -40,21 +40,20 @@ def set_shell_access_creds(config: ET.Element, protocol: Literal['SSH', 'Telnet'
 
 def is_shell_access_enabled(config: ET.Element, interface: Literal['LAN', 'WAN'] = 'LAN',
                             protocol: Literal['SSH', 'Telnet'] = 'SSH') -> bool:
-    ssh_policy = config.find(f".//ServiceControl/*[ServiceList='{protocol.upper()}'][Interface='{interface}']/Policy")
-    return ssh_policy.text == "Permit"
+    service_policy = config.find(f".//ServiceControl/*[ServiceList='{protocol.upper()}'][Interface='{interface}']/Policy")
+    return (service_policy is not None) and (service_policy.text == "Permit")
 
 
 def enable_shell_access(config: ET.Element, interface: Literal['LAN', 'WAN'] = 'LAN',
                         protocol: Literal['SSH', 'Telnet'] = 'SSH', disable: bool = False) -> bool:
-    service_policy = config.find(
-        f".//ServiceControl/*[ServiceList='{protocol.upper()}'][Interface='{interface}']/Policy")
-    remote_management = config.find(f'.//RemoteManagement/{protocol.upper()}/{protocol}Enable')
-    if (service_policy is not None) and (remote_management is not None):
-        service_policy.text = "Discard" if disable else "Permit"
-        if not disable:
-            remote_management.text = '1'
-        return True
-    return False
+    service = config.find(f".//ServiceControl/*[ServiceList='{protocol.upper()}'][Interface='{interface}']")
+    if service is None:
+        return False
+    service_policy = service.find(".Policy")
+    if service_policy is None:
+        service_policy = ET.SubElement(service, 'Policy', attrib={'PARAMETER': 'configured'})
+    service_policy.text = "Discard" if disable else "Permit"
+    return True
 
 
 def create_arg_parser():