ptmanager 0.0.1__py3-none-any.whl

ptmanager/_version.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

ptmanager/modules/config.py

@@ -0,0 +1,95 @@
+from ptlibs import ptprinthelper
+import json
+import os
+import shutil
+
+class Config:
+    NAME = "config.json"
+    PROJECTS_KEY = "projects"
+    TEMP = "temp"
+    SATID_KEY = "satid"
+    PID_KEY = "pid"
+
+    def __init__(self, config_path: str) -> None:
+        self._config_path = config_path
+        self._config: dict[list] = None
+        try:
+            self.load()
+        except (FileNotFoundError, json.JSONDecodeError):
+            ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"Config file not found, creating new ..\n", "ERROR"))
+            self.make()
+
+
+    def load(self) -> dict[list]:
+        with open(self._config_path + self.NAME) as f:
+            self._config = json.load(f)
+
+
+    def make(self) -> dict[list]:
+        self.assure_config_path()
+        with open(self._config_path + self.NAME, "w+") as f:
+            data = {self.SATID_KEY: None, self.PROJECTS_KEY: []}
+            f.write(json.dumps(data, indent=4, sort_keys=True))
+        self._config = json.loads(json.dumps(data))
+
+
+    def assure_config_path(self) -> None:
+        os.makedirs(self._config_path, exist_ok=True)
+
+
+    def delete(self) -> None:
+        os.remove(self._config_path + self.NAME)
+
+    def delete_projects(self) -> None:
+        shutil.rmtree(os.path.join(self._config_path, self.PROJECTS_KEY))
+
+
+    def save(self) -> None:
+        with open(self._config_path + self.NAME, "w") as f:
+            json.dump(self._config, f, indent=4)
+
+
+    def get_path(self) -> str:
+        return self._config_path
+
+    def get_temp_path(self) -> str:
+        temp_path = self._config_path + self.TEMP + "/"
+        os.makedirs(temp_path, exist_ok=True)
+        return temp_path
+
+
+    def print(self) -> None:
+        print("-"*40, json.dumps(self._config, indent=4), "-"*40, sep="\n")
+
+
+    def get_projects(self) -> list:
+        return self._config[self.PROJECTS_KEY]
+
+
+    def get_satid(self) -> str:
+        return self._config[self.SATID_KEY]
+
+
+    def set_satid(self, UID) -> None:
+        self._config[self.SATID_KEY] = UID
+
+
+    def add_project(self, project: dict[str]) -> None:
+        self._config[self.PROJECTS_KEY].append(project)
+
+
+    def get_pid(self, project_id):
+        return self._config[self.PROJECTS_KEY][project_id][self.PID_KEY]
+
+
+    def set_project_pid(self, project_id: int, pid: int) -> None:
+        self._config[self.PROJECTS_KEY][project_id][self.PID_KEY] = pid
+
+
+    def remove_project(self, project_id: int) -> None:
+        shutil.rmtree(os.path.join(self._config_path, self.PROJECTS_KEY, self.get_project(project_id).get("AS-ID")))
+        self._config[self.PROJECTS_KEY].pop(project_id)
+
+
+    def get_project(self, project_id: int):
+        return self._config[self.PROJECTS_KEY][project_id]

ptmanager/modules/process.py

@@ -0,0 +1,25 @@
+import os
+import signal
+
+
+class Process:
+    def __init__(self, PID: int) -> None:
+        self.PID = PID
+
+    def is_running(self) -> bool:
+        "Check if process is running"
+        if self.PID is None:
+            return False
+        try:
+            os.kill(self.PID, 0) # Does not terminate, just checks if running
+            return True
+        except ProcessLookupError:
+            return False
+
+    def kill(self) -> bool:
+        "Tries to kill process with PID"
+        try:
+            os.kill(self.PID, signal.SIGTERM) # Tries to kill process
+            return True
+        except ProcessLookupError:
+            return False

ptmanager/modules/process_manager.py

@@ -0,0 +1,338 @@
+import argparse
+import json; from json.decoder import JSONDecodeError
+import os
+import subprocess
+import sys; sys.path.extend([__file__.rsplit("/", 1)[0], os.path.join(__file__.rsplit("/", 1)[0], "modules")])
+import threading
+import time
+import pathlib
+
+import requests
+
+from process import Process
+from config import Config
+
+class ProcessManager:
+
+    def __init__(self, args):
+        self.config: Config          = Config(f"{str(pathlib.Path.home())}/.ptmanager/")
+        self.satid: str              = self.config.get_satid()
+        self.target: str             = args.target
+        self.API_PATH: str           = "api/v1/sat/"
+
+        self.AS_ID: str              = args.project_id
+        self.project_dir: str        = os.path.join(self.config.get_path(), "projects", self.AS_ID)
+        self.project_tasks_file: str = os.path.join(self.project_dir, "tasks.json")
+
+        self.no_ssl_verify           = args.no_ssl_verify
+        self.proxies: dict           = {"http": args.proxy, "https": args.proxy}
+        self.free_threads            = [i for i in range(args.threads)]
+        self.threads_list            = ["" for _ in range(args.threads)]
+        self.lock                    = threading.Lock()
+
+        if not args.target or not args.auth or not args.sid:
+            self.ptjsonlib.end_error(f"Target, auth and sid are required", self.use_json)
+
+        if not os.path.isdir(self.project_dir):
+            os.makedirs(self.project_dir)
+
+
+    def run(self, args):
+        self.process_front(args.target, args.auth)
+
+
+    def process_front(self, target, auth) -> None:
+        while True:
+            while not self.free_threads:
+                time.sleep(8)
+
+            self.send_results_to_server(target)
+
+            task = self.get_task_from_as(target, auth)
+            if not task:
+                time.sleep(10)
+                continue
+
+            elif task["action"] == "new_task":
+                thread_no = self.free_threads.pop()
+                self.threads_list[thread_no] = threading.Thread(target=self.process_newtask, name=task["guid"], args=(task, thread_no), daemon=False)
+                self.threads_list[thread_no].start()
+            elif task["action"] == "status":
+                self.status_task(task)
+            elif task["action"] == "status-all":
+                self.status_all_tasks()
+            elif task["action"] == "kill-task":
+                self.kill_task(task)
+            elif task["action"] == "kill-all":
+                self.kill_all_tasks()
+            elif task["action"] == "null":
+                pass
+
+
+    def send_results_to_server(self, target) -> None:
+        self.lock.acquire()
+        with self.touchopen(self.project_tasks_file, "r+") as tasks_file:
+            try:
+                tasks_list: list = json.load(tasks_file)
+            except JSONDecodeError:
+                tasks_list = []
+
+
+        for task_idx, task in enumerate(tasks_list):
+            if task["status"] == "running":
+                continue
+            task_result_file = os.path.join(self.project_dir, task["guid"])
+            if os.path.isfile(task_result_file):
+                try:
+                    with open(task_result_file, "r") as file:
+                        task_result = json.load(file)
+                except (JSONDecodeError, Exception) as e:
+                        #TODO pokud existuje odkaz na task v tasks.json, ale neexistuje soubor GUID s taskem, nebo je tento soubor vadný, pak se tento záznam nikdy neodstraní z tasks.json
+                        print("Chyba pri nacitani vysledku automatu ze souboru -", e)
+                        task_result = {}
+                        self.lock.release()
+                        return
+            else:
+                task_result = None
+
+            if task_result:
+                task_result["guid"] = task["guid"]
+                task_result["satid"] = self.satid
+                task_result["results"] = json.dumps(task_result["results"])
+                response = self.send_to_api(end_point="result", data=(task_result))
+                if response.status_code == 200:
+                    # Remove automat result as it's already been posted to AS.
+                    try:
+                        os.remove(task_result_file)
+                    except OSError as e:
+                        pass
+
+                    # TODO: Popnuto z tasks_listu, ted je potreba upraveny (popnuty) tasks_list zaktualizovat v souboru
+                    # TODO: Tak, že otevřu znovu <project_tasks_file> a nahradím jeho obsah popnutym tasks_listem.
+                    tasks_list.pop(task_idx)
+
+                    # Otevře soubor pro zápis - nahradí seznam aktualizovaným tasks_listem.
+                    with self.touchopen(self.project_tasks_file, "w") as tasks_file:
+                        json.dumps(tasks_list, indent=4)
+
+        self.lock.release()
+
+
+    def send_to_api(self, end_point, data) -> requests.Response:
+        target = self.target + self.API_PATH + end_point
+        response = requests.post(target, data=json.dumps(data), proxies=self.proxies, verify=self.no_ssl_verify, headers={"Content-Type": "application/json"})
+        #if response.status_code != 200:
+                #print(f"Response status code is not 200, but {response.status_code}")
+        return response
+
+
+    def status_task(self, task) -> None:
+        """Retrieve status of <task>, repairs tasks.json if task is not running"""
+        self.lock.acquire()
+        with self.touchopen(self.project_tasks_file, "r+") as tasks_list_file:
+            tasks_list = json.loads(tasks_list_file.read())
+            for index, dictionary in enumerate(tasks_list):
+                if dictionary["guid"] == task["guid"]:
+                    if not Process(tasks_list[index]["pid"]).is_running():
+                        tasks_list[index]["status"] = "error"
+                        tasks_list[index]["pid"] = None
+            tasks_list_file.seek(0)
+            tasks_list_file.truncate(0)
+            tasks_list_file.write(json.dumps(tasks_list, indent=4))
+        self.lock.release()
+
+
+    def status_all_tasks(self) -> None:
+        """Repairs all tasks."""
+        self.lock.acquire()
+        with self.touchopen(self.project_tasks_file, "r+") as tasks_list_file:
+            try:
+                tasks_list = json.loads(tasks_list_file.read())
+                for task in tasks_list:
+                    if not Process(task["pid"]).is_running():
+                        task["status"] = "error"
+                        task["pid"] = None
+                tasks_list_file.seek(0)
+                tasks_list_file.truncate(0)
+                tasks_list_file.write(json.dumps(tasks_list, indent=4))
+            except JSONDecodeError:
+                pass
+            finally:
+                self.lock.release()
+
+
+    def kill_all_tasks(self) -> None:
+        """Kills all tasks."""
+
+        for t in self.threads_list:
+            if isinstance(t, threading.Thread):
+                t.join()
+
+        for file in os.listdir(self.project_dir):
+            if file != "tasks.json":
+                os.remove(os.path.join(self.project_dir, file))
+
+        # TODO - KillAllThreads - Pockat nez se thready dokonci, nebo chladnokrevne zavrazdit vsechno co je thread?
+        self.lock.acquire()
+        with self.touchopen(self.project_tasks_file, "r+") as tasks_list_file:
+            try:
+                tasks_list = json.loads(tasks_list_file.read())
+                for task in tasks_list:
+                    if task["pid"]:
+                        Process(task["pid"]).kill()
+                        task["status"] = "killed"
+                        task["pid"] = None
+                tasks_list_file.seek(0)
+                tasks_list_file.truncate(0)
+                tasks_list_file.write(json.dumps(tasks_list, indent=4))
+            except JSONDecodeError:
+                pass
+            finally:
+                self.lock.release()
+
+
+    def kill_task(self, task) -> None:
+        """Kills task with supplied guid."""
+        for t in self.threads_list:
+            if isinstance(t, threading.Thread) and t.name == task["guid"]:
+                t.join()
+        try:
+            os.remove(os.path.join(self.project_dir, task["guid"]))
+        except OSError:
+            # File Not Found
+            pass
+
+        self.lock.acquire()
+        with self.touchopen(self.project_tasks_file, "r+") as tasks_list_file:
+            try:
+                tasks_list = json.loads(tasks_list_file.read())
+                for task_in_list in tasks_list:
+                    if task_in_list["guid"] == task["guid"]:
+                        if task_in_list["pid"]:
+                            Process(task_in_list["pid"]).kill()
+                            task_in_list["status"] = "killed"
+                            task_in_list["pid"] = None
+                tasks_list_file.seek(0)
+                tasks_list_file.truncate(0)
+                tasks_list_file.write(json.dumps(tasks_list, indent=4))
+            except JSONDecodeError:
+                pass
+            finally:
+                self.lock.release()
+
+
+    def touchopen(self, filename, mode):
+        open(filename, "a").close() # "touch" file
+        return open(filename, mode)
+
+    def process_newtask(self, task, thread_no) -> None:
+        """Process newtask with available thread"""
+        result_filename = self.config.get_temp_path() + task["guid"]
+        with open(result_filename, "w+") as result_file:
+            result_file.truncate(0) # Delete file content, if it exist
+            tool_subprocess = subprocess.Popen(task["task"].split(), stdout=result_file, text=True)
+            thread_variable = {"guid": task["guid"], "pid": tool_subprocess.pid, "timeStamp": time.time(), "status": "running"}
+
+        self.lock.acquire()
+        # Read tasks_list_file
+        with self.touchopen(self.project_tasks_file, "r") as tasks_list_file:
+            try:
+                tasks_list = json.load(tasks_list_file)
+            except JSONDecodeError:
+                tasks_list = []
+
+        # Update tasks_list in memory
+        tasks_list.append(thread_variable)
+
+        # Replace content with updated one
+        with self.touchopen(self.project_tasks_file, "w") as tasks_list_file:
+            tasks_list_file.write(json.dumps(tasks_list, indent=4))
+
+        self.lock.release()
+
+        tool_subprocess.wait()
+
+        thread_variable.update({"status": "finished", "pid": None})
+        
+        with open(result_filename, "r") as result_file:
+            tool_result = result_file.read()
+
+        os.remove(result_filename)
+
+
+
+        self.lock.acquire()
+        # Read tasks_list_file
+        with self.touchopen(self.project_tasks_file, "r") as tasks_list_file:
+            tasks_list = json.load(tasks_list_file)
+        self.lock.release()
+
+        for task_idx, task_dict in enumerate(tasks_list):
+            if task_dict["guid"] == thread_variable["guid"]:
+                tasks_list[task_idx] = thread_variable
+
+        self.lock.acquire()
+        # Replace content with updated one
+        with self.touchopen(self.project_tasks_file, "w") as tasks_list_file:
+            tasks_list_file.write(json.dumps(tasks_list, indent=4))
+        self.lock.release()
+
+
+        self.lock.acquire()
+        with open(os.path.join(self.project_dir, task["guid"]), "w") as task_result_file:
+            task_result_file.write(tool_result)
+
+        self.lock.release()
+        self.free_threads.append(thread_no)
+
+
+    def get_task_from_as(self, target=None, auth=None) -> dict:
+        """Retrieve task from AS"""
+        try:
+            response = self.send_to_api(end_point="tasks", data={"satid": self.satid}).json()
+        except requests.RequestException as e:
+            print(f"Chyba pri ziskavani tasku z AS - {e}")
+            return
+
+        try:
+            guid = response["data"]["guid"]
+            action = response["data"]["action"]
+            command = response["data"]["command"]
+            return {"guid": guid, "action": action, "task": command}
+        except Exception as e:
+            return None
+
+
+
+    def _delete_task_from_tasks(self, task) -> None:
+        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), self.AS_ID, "tasks.json"), "r+") as f:
+            original_json = json.loads(f.read())
+            modified_json = [i for i in original_json if i["guid"] != task["guid"]]
+            self.write_to_file_from_start(f, str(modified_json))
+
+
+    def write_to_file_from_start(self, open_file, data: any) -> None:
+        open_file.seek(0)
+        open_file.truncate(0)
+        open_file.write(data)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("-T",    "--target",      type=str)
+    parser.add_argument("-a",    "--auth",        type=str)
+    parser.add_argument("-S",    "--sid",         type=str)
+    parser.add_argument("-prj",  "--project_id",  type=str)
+    parser.add_argument("-p",    "--proxy",       type=str)
+    parser.add_argument("--no_ssl_verify",        action="store_false")
+
+    parser.add_argument("-t",    "--threads",     type=int, default=20)
+
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    tool_subprocess = ProcessManager(args)
+    tool_subprocess.run(args)

ptmanager/modules/project_manager.py

@@ -0,0 +1,195 @@
+import os
+import random
+import signal
+import string
+import subprocess
+import sys; sys.path.extend([__file__.rsplit("/", 1)[0], os.path.join(__file__.rsplit("/", 1)[0], "modules")])
+import urllib
+import uuid
+import json
+
+import requests
+
+from ptlibs import ptjsonlib, ptprinthelper
+
+from config import Config
+from process import Process
+
+
+class ProjectManager:
+    def __init__(self, ptjsonlib: ptjsonlib.PtJsonLib, use_json: bool, proxies: dict, no_ssl_verify: bool, config: Config) -> None:
+        self.ptjsonlib     = ptjsonlib
+        self.use_json      = use_json
+        self.no_ssl_verify = no_ssl_verify
+        self.proxies       = {"http": proxies, "https": proxies}
+        self.config        = config
+
+    def register_project(self, target: str, auth_token: str) -> None:
+        """Registers new project"""
+        if not target:
+            self.ptjsonlib.end_error("Missing --target parameter", self.use_json)
+        if not auth_token:
+            self.ptjsonlib.end_error("Missing --auth parameter", self.use_json)
+
+        target = self.is_url(target)
+        if not target:
+            self.ptjsonlib.end_error("Target is not a valid URL", self.use_json)
+
+        registration_url = target + "api/v1/sat/register"
+        try:
+            response = requests.post(registration_url, proxies=self.proxies, verify=self.no_ssl_verify, data={"token": auth_token, "satid": self.config.get_satid()})
+        except requests.RequestException as e:
+            self.ptjsonlib.end_error("Server is not responding", self.use_json)
+
+        if response.status_code == 200:
+            try:
+                response_json = response.json()
+            except json.JSONDecodeError:
+                self.ptjsonlib.end_error(f"Could not parse target response as json - {response.text}", self.use_json)
+
+            if response_json.get("success"):
+                print(response_json["data"]['name'])
+                print(response_json['message'])
+                project_name = response_json['data']['name']
+            else:
+                self.ptjsonlib.end_error(response_json, self.use_json)
+        else:
+            self.ptjsonlib.end_error(f"[{response.status_code}] {response.text}", self.use_json)
+
+        AS_ID = ''.join(random.choice(string.ascii_lowercase) for _ in range(8))
+        self.config.add_project({"project_name": project_name, "target": target, "auth": auth_token, "pid": None, "AS-ID": AS_ID})
+
+    def is_url(self, url: int):
+        try:
+            result = urllib.parse.urlparse(url)
+            if result.path:
+                while result.path.endswith("/"):
+                    result = result._replace(path=result.path[:-1])
+            result = result._replace(path=result.path + "/")
+            return result.geturl()
+        except ValueError:
+            return False
+
+    def start_project(self, project_id: int) -> None:
+        """Starts specified project."""
+        print("Starting ....")
+        project = self.config.get_project(project_id)
+        if project["pid"]:
+            if not Process(project["pid"]).is_running():
+                self.config.set_project_pid(project_id, None)
+            else:
+                self.ptjsonlib.end_error(f"Project is already running with PID {project['pid']}", self.use_json)
+
+        response = {"sessionid": "sessionid"}
+        try:
+            subprocess_args = [sys.executable, os.path.realpath(os.path.join(__file__.rsplit("/", 1)[0], "process_manager.py")), "--target", project["target"], "--auth", project["auth"], "--sid", response["sessionid"], "--project_id", project["AS-ID"]]
+
+            if self.proxies.get("http"):
+                subprocess_args += ["--proxy", self.proxies.get("http")]
+            if not self.no_ssl_verify:
+                subprocess_args += ["--no_ssl_verify"]
+
+            process = subprocess.Popen(subprocess_args)
+
+        except Exception as e:
+            self.ptjsonlib.end_error(e, self.use_json)
+        self.config.set_project_pid(project_id, process.pid)
+        ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"Started Project {project_id+1} with PID {process.pid}", "INFO"))
+
+
+    def end_project(self, project_id: int) -> None:
+        process_pid = self.config.get_pid(project_id)
+        if process_pid:
+            try:
+                os.kill(process_pid, signal.SIGTERM)
+                self.config.set_project_pid(project_id, None)
+                ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"Killed process with PID {process_pid}", "OK"))
+            except ProcessLookupError:
+                ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"Proccess [{process_pid}] is not running ", "ERROR"))
+                self.config.set_project_pid(project_id, None)
+            except Exception as e:
+                self.ptjsonlib.end_error(e, self.use_json)
+        else:
+            self.ptjsonlib.end_error(f"Project is not running", self.use_json)
+
+    def reset_project(self, project_id: int) -> None:
+        if self.config.get_pid(project_id):
+            self.end_project(project_id)
+            self.start_project(project_id)
+        else:
+            self.start_project(project_id)
+
+        ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"Succesfully registered project", "OK"))
+
+
+    def delete_project(self, project_id: int) -> None:
+        # TODO: Send request to delete project from AS
+        if self.config.get_pid(project_id):
+            self.ptjsonlib.end_error(f"Project is running, end project first", self.use_json)
+
+        project = self.config.get_project(project_id)
+        request_path = project["target"] + "api/v1/sat/delete"
+
+        try:
+            response = requests.post(request_path, proxies=self.proxies, verify=self.no_ssl_verify, data={"satid": self.config.get_satid()}) #project["auth"]
+            self.config.remove_project(project_id)
+            ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"Project deleted succesfully", "OK"))
+        except requests.RequestException as e:
+            ptprinthelper.ptprint(f"Server is not responding", "ERROR")
+            if self._yes_no_prompt("Cannot delete project from AS. Delete TS from server manualy.\nProject will be deleted locally only. Are you sure?", bullet_type=None):
+                self.config.remove_project(project_id)
+                ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"Local project deleted succesfully", "OK"))
+
+
+    def list_projects(self) -> None:
+        print(f"{ptprinthelper.get_colored_text('ID', 'TITLE')}{' '*4}{ptprinthelper.get_colored_text('Project Name', 'TITLE')}{' '*20}{ptprinthelper.get_colored_text('PID', 'TITLE')}{' '*7}{ptprinthelper.get_colored_text('Status', 'TITLE')}{' '*9}")
+        print(f"{'-'*6}{'-'*32}{'-'*10}{'-'*15}")
+        if not self.config.get_projects():
+            print(" ")
+            self.ptjsonlib.end_error("No projects found, register a project first", self.use_json)
+
+        for index, project in enumerate(self.config.get_projects(), 1):
+            if project["pid"]:
+                if not Process(project["pid"]).is_running():
+                    self.config.set_project_pid(index - 1, None)
+                    project["pid"] = None
+            pid = project["pid"]
+            if pid:
+                status = "running"
+            if not pid:
+                status = "-"
+                pid = "-"
+
+            print(f"{index}{' '*(6-len(str(index)))}", end="")
+            print(f"{project['project_name']}{' '*(32-len(project['project_name']))}", end="")
+            print(f"{str(pid)}{' '*(10-len(str(pid)))}", end="")
+            print(f"{status}{' '*(15-len(status))}", end="")
+            print("")
+
+    def register_uid(self) -> None:
+        UID = str(uuid.uuid1())
+        if self.config.get_satid():
+            if self._yes_no_prompt("UID already exists, are you sure you want to create new? All your projects will be deleted!"):
+                self.config.delete_projects()
+                self.config.delete()
+                self.config.make()
+                self.config.set_satid(UID)
+            else:
+                exit()
+        else:
+            self.config.set_satid(UID)
+            print("[*] New UID generated")
+            self.config.print()
+
+
+    def _yes_no_prompt(self, msg, bullet_type="OK") -> bool:
+        ptprinthelper.ptprint_(ptprinthelper.out_ifnot(f"{msg}", bullet_type), end=" ")
+        reply = input(f"y/N: ").upper().strip()
+        if reply == "Y":
+            return True
+        elif reply == "N" or reply == "":
+            return False
+        else:
+            return self._yes_no_prompt(msg)
+
+#TODO: Implementovat přepínač pro insecure SSL a upravit všechny requesty, aby s tímto přepínačem spolupracovaly