paulmataruso/tacticalrmm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

more enum

Browse Source
This commit is contained in:
wh1te909
2022-04-29 06:21:48 +00:00
parent f835997f49
commit 1185ac58e1
8 changed files with 76 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
api/tacticalrmm/agents/management/commands/fake_agents.py
View File

@@ -24,6 +24,7 @@ from tacticalrmm.constants import (
EvtLogNames,
EvtLogTypes,
EvtLogFailWhen,
ScriptShell,
)
from tacticalrmm.demo_data import (
check_network_loc_aware_ps1,
@@ -245,7 +246,7 @@ class Command(BaseCommand):
clear_spool.name = "Clear Print Spooler"
clear_spool.description = "clears the print spooler. Fuck printers"
clear_spool.filename = "clear_print_spool.bat"
clear_spool.shell = "cmd"
clear_spool.shell = ScriptShell.CMD
clear_spool.script_body = clear_print_spool_bat
clear_spool.save()
@@ -253,7 +254,7 @@ class Command(BaseCommand):
check_net_aware.name = "Check Network Location Awareness"
check_net_aware.description = "Check's network location awareness on domain computers, should always be domain profile and not public or private. Sometimes happens when computer restarts before domain available. This script will return 0 if check passes or 1 if it fails."
check_net_aware.filename = "check_network_loc_aware.ps1"
check_net_aware.shell = "powershell"
check_net_aware.shell = ScriptShell.POWERSHELL
check_net_aware.script_body = check_network_loc_aware_ps1
check_net_aware.save()
@@ -261,7 +262,7 @@ class Command(BaseCommand):
check_pool_health.name = "Check storage spool health"
check_pool_health.description = "loops through all storage pools and will fail if any of them are not healthy"
check_pool_health.filename = "check_storage_pool_health.ps1"
check_pool_health.shell = "powershell"
check_pool_health.shell = ScriptShell.POWERSHELL
check_pool_health.script_body = check_storage_pool_health_ps1
check_pool_health.save()
@@ -269,7 +270,7 @@ class Command(BaseCommand):
restart_nla.name = "Restart NLA Service"
restart_nla.description = "restarts the Network Location Awareness windows service to fix the nic profile. Run this after the check network service fails"
restart_nla.filename = "restart_nla.ps1"
restart_nla.shell = "powershell"
restart_nla.shell = ScriptShell.POWERSHELL
restart_nla.script_body = restart_nla_ps1
restart_nla.save()
@@ -277,7 +278,7 @@ class Command(BaseCommand):
show_tmp_dir_script.name = "Check temp dir"
show_tmp_dir_script.description = "shows files in temp dir using python"
show_tmp_dir_script.filename = "show_temp_dir.py"
show_tmp_dir_script.shell = "python"
show_tmp_dir_script.shell = ScriptShell.PYTHON
show_tmp_dir_script.script_body = show_temp_dir_py
show_tmp_dir_script.save()

5
api/tacticalrmm/core/management/commands/post_update_tasks.py
View File

@@ -4,11 +4,10 @@ from django.core.management.base import BaseCommand
from accounts.models import User
from agents.models import Agent
from alerts.models import Alert
from autotasks.models import AutomatedTask
from checks.models import Check, CheckHistory
from scripts.models import Script
from tacticalrmm.constants import AGENT_DEFER
from tacticalrmm.constants import AGENT_DEFER, ScriptType
class Command(BaseCommand):
@@ -27,7 +26,7 @@ class Command(BaseCommand):
user.save()
# convert script base64 field to text field
user_scripts = Script.objects.exclude(script_type="builtin").filter(
user_scripts = Script.objects.exclude(script_type=ScriptType.BUILT_IN).filter(
script_body=""
)
for script in user_scripts:

5
api/tacticalrmm/scripts/baker_recipes.py
View File

@@ -1,9 +1,10 @@
from model_bakery.recipe import Recipe
from tacticalrmm.constants import ScriptShell, ScriptType
script = Recipe(
"scripts.Script",
name="Test Script",
description="Test Desc",
shell="cmd",
script_type="userdefined",
shell=ScriptShell.CMD,
script_type=ScriptType.USER_DEFINED,
)

29
api/tacticalrmm/scripts/models.py
View File

@@ -9,18 +9,7 @@ from django.db.models.fields import CharField, TextField
from logs.models import BaseAuditModel
from tacticalrmm.utils import replace_db_values
SCRIPT_SHELLS = [
("powershell", "Powershell"),
("cmd", "Batch (CMD)"),
("python", "Python"),
("shell", "Shell"),
]
SCRIPT_TYPES = [
("userdefined", "User Defined"),
("builtin", "Built In"),
]
from tacticalrmm.constants import ScriptShell, ScriptType
class Script(BaseAuditModel):
@@ -29,10 +18,10 @@ class Script(BaseAuditModel):
description = models.TextField(null=True, blank=True, default="")
filename = models.CharField(max_length=255, null=True, blank=True)
shell = models.CharField(
max_length=100, choices=SCRIPT_SHELLS, default="powershell"
max_length=100, choices=ScriptShell.choices, default=ScriptShell.POWERSHELL
)
script_type = models.CharField(
max_length=100, choices=SCRIPT_TYPES, default="userdefined"
max_length=100, choices=ScriptType.choices, default=ScriptType.USER_DEFINED
)
args = ArrayField(
models.TextField(null=True, blank=True),
@@ -108,7 +97,9 @@ class Script(BaseAuditModel):
for script in info:
if os.path.exists(os.path.join(scripts_dir, script["filename"])):
s = cls.objects.filter(script_type="builtin", guid=script["guid"])
s = cls.objects.filter(
script_type=ScriptType.BUILT_IN, guid=script["guid"]
)
category = (
script["category"] if "category" in script.keys() else "Community"
@@ -163,7 +154,7 @@ class Script(BaseAuditModel):
name=script["name"],
description=script["description"],
shell=script["shell"],
script_type="builtin",
script_type=ScriptType.BUILT_IN,
category=category,
default_timeout=default_timeout,
args=args,
@@ -178,7 +169,7 @@ class Script(BaseAuditModel):
# check for community scripts that were deleted from json and scripts folder
count, _ = (
Script.objects.filter(script_type="builtin")
Script.objects.filter(script_type=ScriptType.BUILT_IN)
.exclude(guid__in=community_scripts_processed)
.delete()
)
@@ -228,7 +219,9 @@ class ScriptSnippet(models.Model):
name = CharField(max_length=40, unique=True)
desc = CharField(max_length=50, blank=True, default="")
code = TextField(default="")
shell = CharField(max_length=15, choices=SCRIPT_SHELLS, default="powershell")
shell = CharField(
max_length=15, choices=ScriptShell.choices, default=ScriptShell.POWERSHELL
)
def __str__(self):
return self.name

57
api/tacticalrmm/scripts/tests.py
View File

@@ -4,6 +4,7 @@ from django.test import override_settings
from model_bakery import baker
from tacticalrmm.test import TacticalTestCase
from tacticalrmm.constants import ScriptType, ScriptShell
from .models import Script, ScriptSnippet
from .serializers import (
@@ -36,7 +37,7 @@ class TestScriptViews(TacticalTestCase):
data = {
"name": "Name",
"description": "Description",
"shell": "powershell",
"shell": ScriptShell.POWERSHELL,
"category": "New",
"script_body": "Test Script",
"default_timeout": 99,
@@ -93,7 +94,9 @@ class TestScriptViews(TacticalTestCase):
"description": "New Desc",
"script_body": "aasdfdsf",
} # Test
builtin_script = baker.make_recipe("scripts.script", script_type="builtin")
builtin_script = baker.make_recipe(
"scripts.script", script_type=ScriptType.BUILT_IN
)
resp = self.client.put(f"/scripts/{builtin_script.pk}/", data, format="json")
self.assertEqual(resp.status_code, 400)
@@ -136,7 +139,7 @@ class TestScriptViews(TacticalTestCase):
"code": "some_code",
"timeout": 90,
"args": [],
"shell": "powershell",
"shell": ScriptShell.POWERSHELL,
}
resp = self.client.post(url, data, format="json")
@@ -159,7 +162,7 @@ class TestScriptViews(TacticalTestCase):
self.assertFalse(Script.objects.filter(pk=script.pk).exists())
# test delete community script
script = baker.make_recipe("scripts.script", script_type="builtin")
script = baker.make_recipe("scripts.script", script_type=ScriptType.BUILT_IN)
url = f"/scripts/{script.pk}/"
resp = self.client.delete(url, format="json")
self.assertEqual(resp.status_code, 400)
@@ -175,7 +178,9 @@ class TestScriptViews(TacticalTestCase):
# test powershell file
script = baker.make(
"scripts.Script", script_body="Test Script Body", shell="powershell"
"scripts.Script",
script_body="Test Script Body",
shell=ScriptShell.POWERSHELL,
)
url = f"/scripts/{script.pk}/download/"
@@ -187,7 +192,7 @@ class TestScriptViews(TacticalTestCase):
# test batch file
script = baker.make(
"scripts.Script", script_body="Test Script Body", shell="cmd"
"scripts.Script", script_body="Test Script Body", shell=ScriptShell.CMD
)
url = f"/scripts/{script.pk}/download/"
@@ -199,7 +204,7 @@ class TestScriptViews(TacticalTestCase):
# test python file
script = baker.make(
"scripts.Script", script_body="Test Script Body", shell="python"
"scripts.Script", script_body="Test Script Body", shell=ScriptShell.PYTHON
)
url = f"/scripts/{script.pk}/download/"
@@ -228,7 +233,7 @@ class TestScriptViews(TacticalTestCase):
f"-Client '{agent.client.name}'",
f"-Site '{agent.site.name}'",
],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
def test_script_arg_replacement_custom_field(self):
@@ -246,7 +251,7 @@ class TestScriptViews(TacticalTestCase):
# test default value
self.assertEqual(
["-Parameter", "-Another 'DEFAULT'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test with set value
@@ -258,7 +263,7 @@ class TestScriptViews(TacticalTestCase):
)
self.assertEqual(
["-Parameter", "-Another 'CUSTOM VALUE'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
def test_script_arg_replacement_client_custom_fields(self):
@@ -276,7 +281,7 @@ class TestScriptViews(TacticalTestCase):
# test default value
self.assertEqual(
["-Parameter", "-Another 'DEFAULT'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test with set value
@@ -288,7 +293,7 @@ class TestScriptViews(TacticalTestCase):
)
self.assertEqual(
["-Parameter", "-Another 'CUSTOM VALUE'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
def test_script_arg_replacement_site_custom_fields(self):
@@ -306,7 +311,7 @@ class TestScriptViews(TacticalTestCase):
# test default value
self.assertEqual(
["-Parameter", "-Another 'DEFAULT'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test with set value
@@ -318,7 +323,7 @@ class TestScriptViews(TacticalTestCase):
)
self.assertEqual(
["-Parameter", "-Another 'CUSTOM VALUE'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test with set but empty field value
@@ -327,7 +332,7 @@ class TestScriptViews(TacticalTestCase):
self.assertEqual(
["-Parameter", "-Another 'DEFAULT'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test blank default and value
@@ -336,7 +341,7 @@ class TestScriptViews(TacticalTestCase):
self.assertEqual(
["-Parameter", "-Another ''"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
def test_script_arg_replacement_array_fields(self):
@@ -354,7 +359,7 @@ class TestScriptViews(TacticalTestCase):
# test default value
self.assertEqual(
["-Parameter", "-Another 'this,is,an,array'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test with set value and python shell
@@ -366,7 +371,7 @@ class TestScriptViews(TacticalTestCase):
)
self.assertEqual(
["-Parameter", "-Another 'this,is,new'"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
def test_script_arg_replacement_boolean_fields(self):
@@ -384,7 +389,7 @@ class TestScriptViews(TacticalTestCase):
# test default value with python
self.assertEqual(
["-Parameter", "-Another 1"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test with set value and python shell
@@ -396,19 +401,21 @@ class TestScriptViews(TacticalTestCase):
)
self.assertEqual(
["-Parameter", "-Another 0"],
Script.parse_script_args(agent=agent, shell="python", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.PYTHON, args=args),
)
# test with set value and cmd shell
self.assertEqual(
["-Parameter", "-Another 0"],
Script.parse_script_args(agent=agent, shell="cmd", args=args),
Script.parse_script_args(agent=agent, shell=ScriptShell.CMD, args=args),
)
# test with set value and powershell
self.assertEqual(
["-Parameter", "-Another $False"],
Script.parse_script_args(agent=agent, shell="powershell", args=args),
Script.parse_script_args(
agent=agent, shell=ScriptShell.POWERSHELL, args=args
),
)
# test with True value powershell
@@ -417,7 +424,9 @@ class TestScriptViews(TacticalTestCase):
self.assertEqual(
["-Parameter", "-Another $True"],
Script.parse_script_args(agent=agent, shell="powershell", args=args),
Script.parse_script_args(
agent=agent, shell=ScriptShell.POWERSHELL, args=args
),
)
@@ -443,7 +452,7 @@ class TestScriptSnippetViews(TacticalTestCase):
data = {
"name": "Name",
"description": "Description",
"shell": "powershell",
"shell": ScriptShell.POWERSHELL,
"code": "Test",
}

11
api/tacticalrmm/scripts/views.py
View File

@@ -8,6 +8,7 @@ from rest_framework.views import APIView
from agents.permissions import RunScriptPerms
from tacticalrmm.helpers import notify_error
from tacticalrmm.constants import ScriptType, ScriptShell
from .models import Script, ScriptSnippet
from .permissions import ScriptsPerms
@@ -27,7 +28,7 @@ class GetAddScripts(APIView):
showHiddenScripts = request.GET.get("showHiddenScripts", False)
if not showCommunityScripts or showCommunityScripts == "false":
scripts = Script.objects.filter(script_type="userdefined")
scripts = Script.objects.filter(script_type=ScriptType.USER_DEFINED)
else:
scripts = Script.objects.all()
@@ -61,7 +62,7 @@ class GetUpdateDeleteScript(APIView):
data = request.data
if script.script_type == "builtin":
if script.script_type == ScriptType.BUILT_IN:
# allow only favoriting builtin scripts
if "favorite" in data:
# overwrite request data
@@ -83,7 +84,7 @@ class GetUpdateDeleteScript(APIView):
script = get_object_or_404(Script, pk=pk)
# this will never trigger but check anyway
if script.script_type == "builtin":
if script.script_type == ScriptType.BUILT_IN:
return notify_error("Community scripts cannot be deleted")
script.delete()
@@ -172,9 +173,9 @@ def download(request, pk):
if with_snippets == "false":
with_snippets = False
if script.shell == "powershell":
if script.shell == ScriptShell.POWERSHELL:
filename = f"{script.name}.ps1"
elif script.shell == "cmd":
elif script.shell == ScriptShell.CMD:
filename = f"{script.name}.bat"
else:
filename = f"{script.name}.py"

12
api/tacticalrmm/tacticalrmm/constants.py
View File

@@ -19,6 +19,18 @@ CORESETTINGS_CACHE_KEY = "core_settings"
ROLE_CACHE_PREFIX = "role_"
class ScriptShell(models.TextChoices):
POWERSHELL = "powershell", "Powershell"
CMD = "cmd", "Batch (CMD)"
PYTHON = "python", "Python"
SHELL = "shell", "Shell"
class ScriptType(models.TextChoices):
USER_DEFINED = "userdefined", "User Defined"
BUILT_IN = "builtin", "Built In"
class EvtLogNames(models.TextChoices):
APPLICATION = "Application", "Application"
SYSTEM = "System", "System"

6
api/tacticalrmm/tacticalrmm/utils.py
View File

@@ -19,7 +19,7 @@ from agents.models import Agent
from core.models import CodeSignToken
from core.utils import get_core_settings
from logs.models import DebugLog
from tacticalrmm.constants import MONTH_DAYS, MONTHS, WEEK_DAYS, WEEKS
from tacticalrmm.constants import MONTH_DAYS, MONTHS, WEEK_DAYS, WEEKS, ScriptShell
from tacticalrmm.helpers import notify_error
@@ -373,7 +373,7 @@ def replace_db_values(
return ""
def format_shell_array(value: list) -> str:
def format_shell_array(value: list[str]) -> str:
temp_string = ""
for item in value:
temp_string += item + ","
@@ -381,7 +381,7 @@ def format_shell_array(value: list) -> str:
def format_shell_bool(value: bool, shell: Optional[str]) -> str:
if shell == "powershell":
if shell == ScriptShell.POWERSHELL:
return "$True" if value else "$False"
else:
return "1" if value else "0"