d3-tb-agent-comms 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,11 @@
1
+ Metadata-Version: 2.3
2
+ Name: d3-tb-agent-comms
3
+ Version: 0.1.0
4
+ Summary: Lib that holds a Class for easily interacting with TB-Agent Endpoints
5
+ Author: Oscar Mitchell
6
+ Author-email: Oscar Mitchell <oscar.mitchell@disguise.one>
7
+ Requires-Dist: requests>=2.32.5
8
+ Requires-Python: >=3.12
9
+ Description-Content-Type: text/markdown
10
+
11
+ This is a readme
@@ -0,0 +1 @@
1
+ This is a readme
@@ -0,0 +1,21 @@
1
+ [project]
2
+ name = "d3_tb_agent_comms"
3
+ version = "0.1.0"
4
+ description = "Lib that holds a Class for easily interacting with TB-Agent Endpoints"
5
+ readme = "README.md"
6
+ authors = [
7
+ { name = "Oscar Mitchell", email = "oscar.mitchell@disguise.one" }
8
+ ]
9
+ requires-python = ">=3.12"
10
+ dependencies = [
11
+ "requests>=2.32.5",
12
+ ]
13
+
14
+ [build-system]
15
+ requires = ["uv_build>=0.8.15,<0.9.0"]
16
+ build-backend = "uv_build"
17
+
18
+ [dependency-groups]
19
+ dev = [
20
+ "ruff>=0.12.12",
21
+ ]
@@ -0,0 +1,4 @@
1
+ from main import AgentHandler as AgentHandler
2
+ from custom_types import D3SystemInfo as D3SystemInfo, MachineHealthInfo as MachineHealthInfo, AgentHandlerException as AgentHandlerException
3
+
4
+ __all__ = ["AgentHandler", "D3SystemInfo", "MachineHealthInfo", "AgentHandlerException"]
@@ -0,0 +1,27 @@
1
+ from dataclasses import dataclass
2
+ from typing import Final
3
+
4
+
5
+ @dataclass(slots=True)
6
+ class D3SystemInfo(Final):
7
+ major: int
8
+ minor: int
9
+ patch: int
10
+ build: int
11
+ branch: str
12
+ api_port: int
13
+
14
+
15
+ @dataclass(slots=True)
16
+ class MachineHealthInfo(Final):
17
+ revision_number: int
18
+ project_folder_check: bool
19
+ project_folder_shared: bool
20
+ is_service_running: bool
21
+ is_buddy_running: bool
22
+ can_see_storage_server: bool
23
+ can_see_teamcity: bool
24
+
25
+
26
+ class AgentHandlerException(Exception):
27
+ pass
@@ -0,0 +1,184 @@
1
+ # ruff: noqa: N802 - caps in function names
2
+ # In this case the caps serve a purpose of decerning verb vs noun
3
+ # As in we are not, "getting a ping", we are using a "GET" method called "ping"
4
+ from __future__ import annotations
5
+
6
+ import inspect
7
+ import json
8
+ from typing import TYPE_CHECKING, Final
9
+
10
+ import requests
11
+
12
+ from src.tb_agent_comms.types import AgentHandlerException, D3SystemInfo, MachineHealthInfo
13
+
14
+ if TYPE_CHECKING:
15
+ from logging import Logger
16
+
17
+ SUCCESS: Final[int] = 200
18
+
19
+
20
+ class AgentHandler:
21
+ """Handles Agent API requests. Automatic retry and JSON encode/decode included"""
22
+
23
+ def __init__(self, ip: str, port: int, *, logger: Logger) -> None:
24
+ """CTOR for AgentHandler
25
+
26
+ :param ip: the ip of the target TB-Agent
27
+ :param port: the port of the target TB-Agent
28
+ :param logger: a logger instance to use for logging
29
+ """
30
+ self.ip = ip
31
+ self.port = port
32
+ self.logger = logger
33
+
34
+ self.version = -1
35
+
36
+ def _request_handler(self, data: dict | None = None, timeout: int = 5) -> dict:
37
+ """Helper function for handling all request
38
+
39
+ :param data: the data to send to the endpoint
40
+ :param timeout: the timeout for the endpoint
41
+
42
+ :return: the decoded JSON returned from the endpoint
43
+
44
+ :raises:
45
+ """
46
+ calling_func_name = inspect.stack()[1].function
47
+ split_name = calling_func_name.split("_", 1)
48
+ method, endpoint = split_name[0], split_name[1].replace("_", "-")
49
+ url = f"http://{self.ip}:{self.port}/{endpoint}"
50
+ headers = {"accept": "application/json"}
51
+
52
+ attempts = 0
53
+ for _ in range(5):
54
+ attempts += 1
55
+ self.logger.info(f"Attempting #{attempts} call to {endpoint}")
56
+
57
+ try:
58
+ if method == "GET":
59
+ response = requests.get(url=url, headers=headers, timeout=timeout)
60
+ elif method == "POST":
61
+ response = requests.post(url=url, data=data, headers=headers, timeout=timeout)
62
+ else:
63
+ raise NameError(
64
+ f"AgentHandler: Name of calling function did not fit expected format: {calling_func_name}\nThis is a bug with the library" # noqa: E501
65
+ )
66
+ except requests.exceptions.RequestException as e:
67
+ self.logger.error(f"!!!!! Failed to get a response from TB-Agent due to: {e}")
68
+ continue
69
+
70
+ if not response.ok:
71
+ self.logger.warning(
72
+ f"!!! The response from the TB-Agent API was not ok\nCode: {response.status_code}\nResponse text: {response.text}" # noqa: E501
73
+ )
74
+ continue
75
+
76
+ try:
77
+ return json.loads(response.text)
78
+ except ValueError as e:
79
+ self.logger.error(f"!!!!! Failed to decode response from TB-Agent: {e}\nRaw text: {response.text}")
80
+ continue
81
+
82
+ raise AgentHandlerException(f"!!!!! Maximum number of attempts to get a response from /{endpoint} reached")
83
+
84
+ def RAW_ping(self, timeout: int = 2) -> bool:
85
+ """Raw-er version of `GET_ping` for Agent updater checks
86
+
87
+ :return: whether a code 200 came back from the request
88
+ """
89
+ url = f"http://{self.ip}:{self.port}/ping"
90
+ headers = {"accept": "application/json"}
91
+ try:
92
+ response = requests.get(url, headers=headers, timeout=timeout)
93
+ except requests.exceptions.RequestException:
94
+ return False
95
+ else:
96
+ return response.status_code == SUCCESS
97
+
98
+ def GET_ping(self) -> None:
99
+ return self._request_handler()["return"]
100
+
101
+ def GET_d3installed(self) -> bool:
102
+ return self._request_handler()["d3installed"]
103
+
104
+ def GET_d3system(self) -> D3SystemInfo:
105
+ body = self._request_handler()
106
+ return D3SystemInfo(
107
+ major=body["major"],
108
+ minor=body["minor"],
109
+ patch=body["patch"],
110
+ build=body["build"],
111
+ branch=body["branch"],
112
+ api_port=body["apiPort"],
113
+ )
114
+
115
+ def GET_agent_version(self) -> int:
116
+ return self._request_handler()["version"]
117
+
118
+ def GET_running_project(self) -> str:
119
+ return self._request_handler()["projectName"]
120
+
121
+ def GET_summarise_machine_health(self) -> MachineHealthInfo:
122
+ body = self._request_handler()
123
+ return MachineHealthInfo(
124
+ revision_number=body["revision number"],
125
+ project_folder_check=body["project folder check"],
126
+ project_folder_shared=body["shared project folder check"],
127
+ is_service_running=body["is d3 service running"],
128
+ is_buddy_running=body["is d3 buddy running"],
129
+ can_see_storage_server=body["can machine see storage server"],
130
+ can_see_teamcity=body["can machine see teamcity"],
131
+ )
132
+
133
+ def GET_gpu_vendor(self) -> str:
134
+ return self._request_handler()["gpu_output"]
135
+
136
+ def GET_current_os(self) -> str:
137
+ return self._request_handler()["OS"]
138
+
139
+ def GET_licence_codes(self) -> str:
140
+ return self._request_handler()["licenceCodes"]
141
+
142
+ def POST_restart_local_machine(self, timeout: int = 1) -> str:
143
+ data = {"timeout_duration": timeout}
144
+ return self._request_handler(data)["msg"]
145
+
146
+ def POST_download_d3(
147
+ self, revision_number: int, destination: str, timeout: int = 0, download_time: int = 360, host: str = ""
148
+ ) -> str:
149
+ data = {"rev": revision_number, "destination": destination}
150
+ if timeout:
151
+ data["timeout"] = timeout
152
+ if host:
153
+ data["host"] = host
154
+ return self._request_handler(data, timeout=download_time)["msg"]
155
+
156
+ def POST_install_d3(self, path: str) -> str:
157
+ data = {"path": path}
158
+ return self._request_handler(data, timeout=15)["taskName"]
159
+
160
+ def POST_check_pid(self, pid: int) -> bool:
161
+ data = {"pid": pid}
162
+ return self._request_handler(data)["runnning"]
163
+
164
+ def POST_update_agent(self, github_key: str) -> None:
165
+ data = {"githubKey": github_key}
166
+ self._request_handler(data, timeout=25)
167
+
168
+ def POST_start_d3_project(self, project_folder: str, project_name: str) -> int:
169
+ data = {"projectFolder": project_folder, "projectName": project_name}
170
+ return self._request_handler(data)["msg"]
171
+
172
+ def POST_terminate_d3(self) -> str:
173
+ return self._request_handler({})["msg"]
174
+
175
+ def POST_delete_d3_installer(self, installer_path: str) -> str:
176
+ data = {"path": installer_path}
177
+ return self._request_handler(data)["msg"]
178
+
179
+ def POST_check_task_running(self, task_name: str) -> bool:
180
+ data = {"taskName": task_name}
181
+ return self._request_handler(data, timeout=15)["isRunning"]
182
+
183
+ def POST_restart_d3service(self) -> str:
184
+ return self._request_handler()["msg"]
File without changes