tacticalrmm/_modules/win_agent.py

458 lines
13 KiB
Python

from __future__ import absolute_import
import psutil
import os
import datetime
import zlib
import json
import base64
import wmi
import win32evtlog
import win32con
import win32evtlogutil
import winerror
from time import sleep
import requests
import subprocess
import random
import platform
ARCH = "64" if platform.machine().endswith("64") else "32"
PROGRAM_DIR = os.path.join(os.environ["ProgramFiles"], "TacticalAgent")
TAC_RMM = os.path.join(PROGRAM_DIR, "tacticalrmm.exe")
NSSM = os.path.join(PROGRAM_DIR, "nssm.exe" if ARCH == "64" else "nssm-x86.exe")
TEMP_DIR = os.path.join(os.environ["WINDIR"], "Temp")
SYS_DRIVE = os.environ["SystemDrive"]
PY_BIN = os.path.join(SYS_DRIVE, "\\salt", "bin", "python.exe")
SALT_CALL = os.path.join(SYS_DRIVE, "\\salt", "salt-call.bat")
def get_services():
# see https://github.com/wh1te909/tacticalrmm/issues/38
# for why I am manually implementing the svc.as_dict() method of psutil
ret = []
for svc in psutil.win_service_iter():
i = {}
try:
i["display_name"] = svc.display_name()
i["binpath"] = svc.binpath()
i["username"] = svc.username()
i["start_type"] = svc.start_type()
i["status"] = svc.status()
i["pid"] = svc.pid()
i["name"] = svc.name()
i["description"] = svc.description()
except Exception:
continue
else:
ret.append(i)
return ret
def run_python_script(filename, timeout, script_type="userdefined"):
# no longer used in agent version 0.11.0
file_path = os.path.join(TEMP_DIR, filename)
if os.path.exists(file_path):
try:
os.remove(file_path)
except:
pass
if script_type == "userdefined":
__salt__["cp.get_file"](f"salt://scripts/userdefined/{filename}", file_path)
else:
__salt__["cp.get_file"](f"salt://scripts/{filename}", file_path)
return __salt__["cmd.run_all"](f"{PY_BIN} {file_path}", timeout=timeout)
def run_script(filepath, filename, shell, timeout, args=[], bg=False):
if shell == "powershell" or shell == "cmd":
if args:
return __salt__["cmd.script"](
source=filepath,
args=" ".join(map(lambda x: f'"{x}"', args)),
shell=shell,
timeout=timeout,
bg=bg,
)
else:
return __salt__["cmd.script"](
source=filepath, shell=shell, timeout=timeout, bg=bg
)
elif shell == "python":
file_path = os.path.join(TEMP_DIR, filename)
if os.path.exists(file_path):
try:
os.remove(file_path)
except:
pass
__salt__["cp.get_file"](filepath, file_path)
salt_cmd = "cmd.run_bg" if bg else "cmd.run_all"
if args:
a = " ".join(map(lambda x: f'"{x}"', args))
cmd = f"{PY_BIN} {file_path} {a}"
return __salt__[salt_cmd](cmd, timeout=timeout)
else:
return __salt__[salt_cmd](f"{PY_BIN} {file_path}", timeout=timeout)
def uninstall_agent():
remove_exe = os.path.join(PROGRAM_DIR, "unins000.exe")
__salt__["cmd.run_bg"]([remove_exe, "/VERYSILENT", "/SUPPRESSMSGBOXES"])
return "ok"
def update_salt():
for p in psutil.process_iter():
with p.oneshot():
if p.name() == "tacticalrmm.exe" and "updatesalt" in p.cmdline():
return "running"
from subprocess import Popen, PIPE
CREATE_NEW_PROCESS_GROUP = 0x00000200
DETACHED_PROCESS = 0x00000008
cmd = [TAC_RMM, "-m", "updatesalt"]
p = Popen(
cmd,
stdin=PIPE,
stdout=PIPE,
stderr=PIPE,
close_fds=True,
creationflags=DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP,
)
return p.pid
def run_manual_checks():
__salt__["cmd.run_bg"]([TAC_RMM, "-m", "runchecks"])
return "ok"
def install_updates():
for p in psutil.process_iter():
with p.oneshot():
if p.name() == "tacticalrmm.exe" and "winupdater" in p.cmdline():
return "running"
return __salt__["cmd.run_bg"]([TAC_RMM, "-m", "winupdater"])
def _wait_for_service(svc, status, retries=10):
attempts = 0
while 1:
try:
service = psutil.win_service_get(svc)
except psutil.NoSuchProcess:
stat = "fail"
attempts += 1
sleep(5)
else:
stat = service.status()
if stat != status:
attempts += 1
sleep(5)
else:
attempts = 0
if attempts == 0 or attempts > retries:
break
return stat
def agent_update_v2(inno, url):
# make sure another instance of the update is not running
# this function spawns 2 instances of itself (because we call it twice with salt run_bg)
# so if more than 2 running, don't continue as an update is already running
count = 0
for p in psutil.process_iter():
try:
with p.oneshot():
if "win_agent.agent_update_v2" in p.cmdline():
count += 1
except Exception:
continue
if count > 2:
return "already running"
sleep(random.randint(1, 20)) # don't flood the rmm
exe = os.path.join(TEMP_DIR, inno)
if os.path.exists(exe):
try:
os.remove(exe)
except:
pass
try:
r = requests.get(url, stream=True, timeout=600)
except Exception:
return "failed"
if r.status_code != 200:
return "failed"
with open(exe, "wb") as f:
for chunk in r.iter_content(chunk_size=1024):
if chunk:
f.write(chunk)
del r
ret = subprocess.run([exe, "/VERYSILENT", "/SUPPRESSMSGBOXES"], timeout=120)
tac = _wait_for_service(svc="tacticalagent", status="running")
if tac != "running":
subprocess.run([NSSM, "start", "tacticalagent"], timeout=30)
chk = _wait_for_service(svc="checkrunner", status="running")
if chk != "running":
subprocess.run([NSSM, "start", "checkrunner"], timeout=30)
return "ok"
def do_agent_update_v2(inno, url):
return __salt__["cmd.run_bg"](
[
SALT_CALL,
"win_agent.agent_update_v2",
f"inno={inno}",
f"url={url}",
"--local",
]
)
def agent_update(version, url):
# make sure another instance of the update is not running
# this function spawns 2 instances of itself so if more than 2 running,
# don't continue as an update is already running
count = 0
for p in psutil.process_iter():
try:
with p.oneshot():
if "win_agent.agent_update" in p.cmdline():
count += 1
except Exception:
continue
if count > 2:
return "already running"
sleep(random.randint(1, 60)) # don't flood the rmm
try:
r = requests.get(url, stream=True, timeout=600)
except Exception:
return "failed"
if r.status_code != 200:
return "failed"
exe = os.path.join(TEMP_DIR, f"winagent-v{version}.exe")
with open(exe, "wb") as f:
for chunk in r.iter_content(chunk_size=1024):
if chunk:
f.write(chunk)
del r
services = ("tacticalagent", "checkrunner")
for svc in services:
subprocess.run([NSSM, "stop", svc], timeout=120)
sleep(10)
r = subprocess.run([exe, "/VERYSILENT", "/SUPPRESSMSGBOXES"], timeout=300)
sleep(30)
for svc in services:
subprocess.run([NSSM, "start", svc], timeout=120)
return "ok"
def do_agent_update(version, url):
return __salt__["cmd.run_bg"](
[
SALT_CALL,
"win_agent.agent_update",
f"version={version}",
f"url={url}",
"--local",
]
)
class SystemDetail:
def __init__(self):
self.c = wmi.WMI()
self.comp_sys_prod = self.c.Win32_ComputerSystemProduct()
self.comp_sys = self.c.Win32_ComputerSystem()
self.memory = self.c.Win32_PhysicalMemory()
self.os = self.c.Win32_OperatingSystem()
self.base_board = self.c.Win32_BaseBoard()
self.bios = self.c.Win32_BIOS()
self.disk = self.c.Win32_DiskDrive()
self.network_adapter = self.c.Win32_NetworkAdapter()
self.network_config = self.c.Win32_NetworkAdapterConfiguration()
self.desktop_monitor = self.c.Win32_DesktopMonitor()
self.cpu = self.c.Win32_Processor()
self.usb = self.c.Win32_USBController()
def get_all(self, obj):
ret = []
for i in obj:
tmp = [
{j: getattr(i, j)}
for j in list(i.properties)
if getattr(i, j) is not None
]
ret.append(tmp)
return ret
def system_info():
info = SystemDetail()
return {
"comp_sys_prod": info.get_all(info.comp_sys_prod),
"comp_sys": info.get_all(info.comp_sys),
"mem": info.get_all(info.memory),
"os": info.get_all(info.os),
"base_board": info.get_all(info.base_board),
"bios": info.get_all(info.bios),
"disk": info.get_all(info.disk),
"network_adapter": info.get_all(info.network_adapter),
"network_config": info.get_all(info.network_config),
"desktop_monitor": info.get_all(info.desktop_monitor),
"cpu": info.get_all(info.cpu),
"usb": info.get_all(info.usb),
}
def local_sys_info():
return __salt__["cmd.run_bg"]([TAC_RMM, "-m", "sysinfo"])
def get_procs():
ret = []
# setup
for proc in psutil.process_iter():
with proc.oneshot():
proc.cpu_percent(interval=None)
# need time for psutil to record cpu percent
sleep(1)
for c, proc in enumerate(psutil.process_iter(), 1):
x = {}
with proc.oneshot():
if proc.pid == 0 or not proc.name():
continue
x["name"] = proc.name()
x["cpu_percent"] = proc.cpu_percent(interval=None) / psutil.cpu_count()
x["memory_percent"] = proc.memory_percent()
x["pid"] = proc.pid
x["ppid"] = proc.ppid()
x["status"] = proc.status()
x["username"] = proc.username()
x["id"] = c
ret.append(x)
return ret
def _compress_json(j):
return {
"wineventlog": base64.b64encode(
zlib.compress(json.dumps(j).encode("utf-8", errors="ignore"))
).decode("ascii", errors="ignore")
}
def get_eventlog(logtype, last_n_days):
start_time = datetime.datetime.now() - datetime.timedelta(days=last_n_days)
flags = win32evtlog.EVENTLOG_BACKWARDS_READ | win32evtlog.EVENTLOG_SEQUENTIAL_READ
status_dict = {
win32con.EVENTLOG_AUDIT_FAILURE: "AUDIT_FAILURE",
win32con.EVENTLOG_AUDIT_SUCCESS: "AUDIT_SUCCESS",
win32con.EVENTLOG_INFORMATION_TYPE: "INFO",
win32con.EVENTLOG_WARNING_TYPE: "WARNING",
win32con.EVENTLOG_ERROR_TYPE: "ERROR",
0: "INFO",
}
computer = "localhost"
hand = win32evtlog.OpenEventLog(computer, logtype)
total = win32evtlog.GetNumberOfEventLogRecords(hand)
log = []
uid = 0
done = False
try:
while 1:
events = win32evtlog.ReadEventLog(hand, flags, 0)
for ev_obj in events:
uid += 1
# return once total number of events reach or we'll be stuck in an infinite loop
if uid >= total:
done = True
break
the_time = ev_obj.TimeGenerated.Format()
time_obj = datetime.datetime.strptime(the_time, "%c")
if time_obj < start_time:
done = True
break
computer = str(ev_obj.ComputerName)
src = str(ev_obj.SourceName)
evt_type = str(status_dict[ev_obj.EventType])
evt_id = str(winerror.HRESULT_CODE(ev_obj.EventID))
evt_category = str(ev_obj.EventCategory)
record = str(ev_obj.RecordNumber)
msg = (
str(win32evtlogutil.SafeFormatMessage(ev_obj, logtype))
.replace("<", "")
.replace(">", "")
)
event_dict = {
"computer": computer,
"source": src,
"eventType": evt_type,
"eventID": evt_id,
"eventCategory": evt_category,
"message": msg,
"time": the_time,
"record": record,
"uid": uid,
}
log.append(event_dict)
if done:
break
except Exception:
pass
win32evtlog.CloseEventLog(hand)
return _compress_json(log)