diff --git a/api/tacticalrmm/agents/tasks.py b/api/tacticalrmm/agents/tasks.py index c3cc03cc..df3a669b 100644 --- a/api/tacticalrmm/agents/tasks.py +++ b/api/tacticalrmm/agents/tasks.py @@ -1,7 +1,6 @@ import asyncio import datetime as dt import random -import requests import urllib.parse from time import sleep from typing import Union @@ -21,21 +20,9 @@ from tacticalrmm.utils import run_nats_api_cmd logger.configure(**settings.LOG_CONFIG) -def _get_exegen_url() -> str: - urls: list[str] = settings.EXE_GEN_URLS - for url in urls: - try: - r = requests.get(url, timeout=10) - except: - continue - - if r.status_code == 200: - return url - - return random.choice(urls) - - def agent_update(pk: int, codesigntoken: str = None) -> str: + from agents.utils import get_exegen_url + agent = Agent.objects.get(pk=pk) if pyver.parse(agent.version) <= pyver.parse("1.3.0"): @@ -52,7 +39,7 @@ def agent_update(pk: int, codesigntoken: str = None) -> str: inno = agent.win_inno_exe if codesigntoken is not None and pyver.parse(version) >= pyver.parse("1.5.0"): - base_url = _get_exegen_url() + "/api/v1/winagents/?" + base_url = get_exegen_url() + "/api/v1/winagents/?" params = {"version": version, "arch": agent.arch, "token": codesigntoken} url = base_url + urllib.parse.urlencode(params) else: diff --git a/api/tacticalrmm/agents/tests.py b/api/tacticalrmm/agents/tests.py index 60b9d2fe..4a0c56c9 100644 --- a/api/tacticalrmm/agents/tests.py +++ b/api/tacticalrmm/agents/tests.py @@ -914,7 +914,7 @@ class TestAgentTasks(TacticalTestCase): self.authenticate() self.setup_coresettings() - @patch("agents.tasks._get_exegen_url") + @patch("agents.utils.get_exegen_url") @patch("agents.models.Agent.nats_cmd") def test_agent_update(self, nats_cmd, get_exe): from agents.tasks import agent_update diff --git a/api/tacticalrmm/agents/utils.py b/api/tacticalrmm/agents/utils.py new file mode 100644 index 00000000..6180fedc --- /dev/null +++ b/api/tacticalrmm/agents/utils.py @@ -0,0 +1,37 @@ +import random +import requests +import urllib.parse + +from django.conf import settings + + +def get_exegen_url() -> str: + urls: list[str] = settings.EXE_GEN_URLS + for url in urls: + try: + r = requests.get(url, timeout=10) + except: + continue + + if r.status_code == 200: + return url + + return random.choice(urls) + + +def get_winagent_url(arch: str) -> str: + from core.models import CodeSignToken + + try: + codetoken = CodeSignToken.objects.first().token + base_url = get_exegen_url() + "/api/v1/winagents/?" + params = { + "version": settings.LATEST_AGENT_VER, + "arch": arch, + "token": codetoken, + } + dl_url = base_url + urllib.parse.urlencode(params) + except: + dl_url = settings.DL_64 if arch == "64" else settings.DL_32 + + return dl_url diff --git a/api/tacticalrmm/agents/views.py b/api/tacticalrmm/agents/views.py index a92b1a62..1c718568 100644 --- a/api/tacticalrmm/agents/views.py +++ b/api/tacticalrmm/agents/views.py @@ -353,6 +353,7 @@ class Reboot(APIView): @api_view(["POST"]) def install_agent(request): + from agents.utils import get_winagent_url from knox.models import AuthToken client_id = request.data["client"] @@ -375,7 +376,7 @@ def install_agent(request): inno = ( f"winagent-v{version}.exe" if arch == "64" else f"winagent-v{version}-x86.exe" ) - download_url = settings.DL_64 if arch == "64" else settings.DL_32 + download_url = get_winagent_url(arch) _, token = AuthToken.objects.create( user=request.user, expiry=dt.timedelta(hours=request.data["expires"]) diff --git a/api/tacticalrmm/tacticalrmm/utils.py b/api/tacticalrmm/tacticalrmm/utils.py index 301348d5..72fea20d 100644 --- a/api/tacticalrmm/tacticalrmm/utils.py +++ b/api/tacticalrmm/tacticalrmm/utils.py @@ -52,7 +52,7 @@ def generate_winagent_exe( file_name: str, ) -> Union[Response, FileResponse]: - from agents.tasks import _get_exegen_url + from agents.utils import get_exegen_url inno = ( f"winagent-v{settings.LATEST_AGENT_VER}.exe" @@ -62,7 +62,7 @@ def generate_winagent_exe( try: codetoken = CodeSignToken.objects.first().token - base_url = _get_exegen_url() + "/api/v1/winagents/?" + base_url = get_exegen_url() + "/api/v1/winagents/?" params = { "version": settings.LATEST_AGENT_VER, "arch": arch,