pyflow-net 0.1.0__py3-none-any.whl

PyFlow/__init__.py

@@ -0,0 +1,4 @@
+from . import flow_setup
+
+def main():
+    flow_setup.main()

PyFlow/__main__.py

@@ -0,0 +1,4 @@
+from . import main
+
+if __name__=="__main__":
+    main()

PyFlow/command_control_extension_tcp.py

@@ -0,0 +1,252 @@
+import os
+import ast
+import json
+import copy
+import shlex
+import shutil
+import traceback
+import threading
+import subprocess
+from datetime import datetime
+from .network_api import connect_tcp
+
+server_instance = None
+client_instance = None
+command_counter = {}
+command_counter_lock = threading.Lock()
+
+
+def _load_json_file(path):
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            return data if isinstance(data, dict) else {}
+    except Exception:
+        return {}
+
+
+def _merge_log_dicts(base_data, extra_data):
+    merged = {}
+    for command, entries in base_data.items():
+        merged[command] = list(entries) if isinstance(entries, list) else [entries]
+    for command, entries in extra_data.items():
+        if command not in merged:
+            merged[command] = []
+        if isinstance(entries, list):
+            merged[command].extend(entries)
+        else:
+            merged[command].append(entries)
+    return merged
+
+
+def _merge_all_logs(log_dir, merged_filename="merged_logs.json"):
+    merged_data = {}
+    for filename in sorted(os.listdir(log_dir)):
+        if not filename.endswith(".json") or filename == merged_filename:
+            continue
+        file_path = os.path.join(log_dir, filename)
+        if os.path.isfile(file_path):
+            merged_data = _merge_log_dicts(merged_data, _load_json_file(file_path))
+    with open(os.path.join(log_dir, merged_filename), "w", encoding="utf-8") as f:
+        json.dump(merged_data, f, ensure_ascii=False, indent=2)
+
+
+def _setup_command():
+    print("Setting up server command...")
+    server_instance.register_command(
+        command_name="/command", handler=_command_handler, where_to_run="client", run_in_thread=True
+    )
+    server_instance.register_command(
+        command_name="/command_done",
+        handler=_command_done_dealing_server,
+        where_to_run="server",
+        run_in_thread=True,
+    )
+
+
+def _setup_client_command():
+    print("Setting up client command...")
+    client_instance.register_command(
+        command_name="/command",
+        handler=_command_handler_server_setup,
+        where_to_run="server",
+        run_in_thread=True,
+    )
+
+
+def _command_handler(sock, addr, cmd):
+    print(f"Received command from {addr}: {cmd}")
+    client_class = server_instance.clients
+    cmd_part = shlex.split(cmd)
+    del cmd_part[0]  # Remove the command name
+    command_client_pair = []
+    clients_list = []
+    commands_list = []
+    clients_num = 0
+    for part in cmd_part:
+        if part.startswith("(") and part.endswith(")"):
+            try:
+                clients_num += 1
+                client_part = ast.literal_eval(part)
+                clients_list.append(client_part)
+                print(f"client part: {clients_list}")
+            except Exception as e:
+                print(f"Error evaluating part '{part}': {e}")
+        else:
+            if clients_num != 0:
+                command_client_pair.append([commands_list, clients_list])
+                clients_num = 0
+                clients_list = []
+                commands_list = []
+            commands_list.append(part)
+            print(f"command part: {commands_list}")
+    if clients_num != 0:
+        print([commands_list, clients_list])
+        command_client_pair.append([commands_list, clients_list])
+        clients_num = 0
+        clients_list = []
+        commands_list = []
+    client_id = 0
+    for pair in command_client_pair:
+        for msg in pair[0]:
+            command_msg = (
+                "/command" + " " + shlex.quote(msg) + " " + shlex.quote(str(len(pair[0]))) + " "
+            )
+            for client in pair[1]:
+                temp_msg = (
+                    command_msg
+                    + shlex.quote(str(client_id))
+                    + " "
+                    + shlex.quote(str(client))
+                    + "\n"
+                )
+                client_socket = client_class[client]["socket"]
+                server_instance.send_message(client_socket=client_socket, message=temp_msg)
+                print(f"Sending command to clients: {command_msg}")
+            client_id += 1
+
+
+def _command_handler_server_setup(sock, addr, cmd):
+    global command_counter
+    print(f"Received command from {addr}: {cmd}")
+    try:
+        cmd_parts = shlex.split(cmd)
+    except Exception as e:
+        print(f"Error parsing command: {e}")
+        return
+    if len(cmd_parts) < 3 or cmd_parts[0] != "/command":
+        print("Invalid command format.")
+        return
+    command = cmd_parts[1]
+    client_addr = cmd_parts[4]
+    command_total_num = int(cmd_parts[2])
+    client_id = cmd_parts[3]
+    log_dir = os.path.join(os.path.dirname(__file__), "logs")
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    with command_counter_lock:
+        if str(client_id) not in command_counter:
+            command_counter[str(client_id)] = 1
+        else:
+            command_counter[str(client_id)] += 1
+        cmd_id = copy.copy(command_counter[str(client_id)])
+    log_filename = "logs{}.json".format("_" + str(client_id))
+    log_path = os.path.join(log_dir, log_filename)
+    try:
+        result = subprocess.run(command, capture_output=True, text=True, shell=True)
+        output = result.stdout.strip()
+        error = result.stderr.strip()
+        returncode = result.returncode
+    except Exception as e:
+        output = ""
+        error = str(e)
+        returncode = -1
+    log_line = {
+        "timestamp": datetime.now().isoformat(),
+        "command": command,
+        "cmd_id": str(cmd_id),
+        "client_id": str(client_id),
+        "client": client_addr,
+        "from": str(addr),
+        "output": output,
+        "error": error,
+        "returncode": returncode,
+    }
+    if os.path.exists(log_path):
+        with open(log_path, "r", encoding="utf-8") as f:
+            try:
+                log_data = json.load(f)
+            except Exception:
+                log_data = {}
+    else:
+        log_data = {}
+    if command not in log_data:
+        log_data[command] = []
+    log_data[command].append(log_line)
+    with open(log_path, "w", encoding="utf-8") as f:
+        json.dump(log_data, f, ensure_ascii=False, indent=2)
+    print(f"Log written to {log_path}")
+    msg = '/file "{}"'.format(log_path)
+    if cmd_id == command_total_num:
+        command_counter[str(client_id)] = 0
+        client_instance.file_transfer_client_recv_client_start(
+            message=msg, file_folder_abspath=None
+        )
+        client_instance.send_message(
+            client_socket=client_instance.client_socket,
+            message='/command_done "{}" "{}"'.format(log_filename, log_path),
+        )
+    else:
+        pass
+    print("Dealing the command successfully!")
+
+
+def _command_done_dealing_server(sock, addr, cmd):
+    cmd_parts = shlex.split(cmd)
+    log_filename = cmd_parts[1]
+    log_path = cmd_parts[2]
+    log_dir = os.path.join(os.path.dirname(__file__), "logs")
+    os.makedirs(log_dir, exist_ok=True)
+    received_log_file = os.path.join(server_instance.file_transfer_dir, os.path.basename(log_path))
+    destination_log_file = os.path.join(log_dir, log_filename)
+    try:
+        if os.path.exists(destination_log_file):
+            existing_logs = _load_json_file(destination_log_file)
+            incoming_logs = _load_json_file(received_log_file)
+            merged_logs = _merge_log_dicts(existing_logs, incoming_logs)
+            with open(destination_log_file, "w", encoding="utf-8") as f:
+                json.dump(merged_logs, f, ensure_ascii=False, indent=2)
+            os.remove(received_log_file)
+        else:
+            shutil.move(received_log_file, destination_log_file)
+        _merge_all_logs(log_dir)
+        print("Command Done!")
+    except Exception:
+        traceback.print_exc()
+        print("ErrorWhileMovingTheLogFile: moving log file failed.")
+
+
+def client_setup():
+    global client_instance
+    client_instance = connect_tcp.TCP_Client_Base(
+        host="127.0.0.1",
+        port=65000,
+        client_host="127.0.0.1",
+        is_input_command_in_console=True,
+        is_extend_command=True,
+    )
+    _setup_client_command()
+    client_instance.start_TCP_client()
+
+
+def server_setup():
+    global server_instance
+    server_instance = connect_tcp.TCP_Server_Base(
+        host="127.0.0.1",
+        port=65000,
+        max_clients=10,
+        is_input_command_in_console=True,
+        is_extend_command=True,
+    )
+    _setup_command()
+    server_instance.start_TCP_Server()

PyFlow/flow_setup.py

@@ -0,0 +1,283 @@
+import os
+import sys
+import json
+import shutil
+import platform
+import tempfile
+import argparse
+import subprocess
+import traceback
+from .network_api.connect_tcp import TCP_Server_Base, TCP_Client_Base
+
+SERVER_DEFAULTS = {
+    "host": "127.0.0.1",
+    "port": 65432,
+    "max_clients": 10,
+    "port_add_step": 1,
+    "port_range_num": 100,
+    "max_file_transfer_thread_num": 10,
+    "is_hand_alloc_port": False,
+    "is_input_command_in_console": True,
+    "max_custom_workers": 10,
+    "is_extend_command": False,
+    "is_enable_encrypto": True,
+    "is_custom_keys": None,
+}
+
+CLIENT_DEFAULTS = {
+    "host": None,
+    "client_host": "127.0.0.1",
+    "port": 65432,
+    "client_port": None,
+    "timeout": None,
+    "port_add_step": 1,
+    "max_thread_num": 10,
+    "is_input_command_in_console": True,
+    "is_wait_server": True,
+    "max_custom_workers": 10,
+    "is_extend_command": False,
+    "is_enable_encrypto": True,
+    "is_custom_keys": None,
+}
+
+
+def parse_addr_port(addr_port):
+    host, port_str = addr_port.strip().split(":")
+    return host, int(port_str)
+
+
+def load_existing_config():
+    if os.path.exists("setup.json"):
+        with open("setup.json", "r", encoding="utf-8") as f:
+            return json.load(f)
+    return {"servers": [], "clients": []}
+
+
+def complete_server_config(cfg):
+    full = SERVER_DEFAULTS.copy()
+    full.update(cfg)
+    return full
+
+
+def complete_client_config(cfg):
+    full = CLIENT_DEFAULTS.copy()
+    full.update(cfg)
+    return full
+
+
+def save_config(servers, clients):
+    if servers:
+        servers = [servers[-1]]
+    else:
+        servers = []
+    if clients:
+        clients = [clients[-1]]
+    else:
+        clients = []
+    data = {
+        "servers": [complete_server_config(cfg) for cfg in servers],
+        "clients": [complete_client_config(cfg) for cfg in clients],
+    }
+    with open("setup.json", "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4, ensure_ascii=False)
+
+
+def launch_instance(config, instance_type):
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as f:
+        json.dump(config, f)
+        config_file_path = f.name
+    script = os.path.abspath(__file__)
+    python = sys.executable
+    launch_arg = f"--launch_{instance_type}"
+    system = platform.system()
+    try:
+        if system == "Windows":
+            cmd = f'start cmd /k {python} {script} {launch_arg} --config-file "{config_file_path}"'
+            subprocess.Popen(cmd, shell=True)
+        elif system == "Linux":
+            terminals = ["gnome-terminal", "xterm", "x-terminal-emulator"]
+            launched = False
+            for term in terminals:
+                if shutil.which(term):
+                    cmd = f'{term} -- {python} {script} {launch_arg} --config-file "{config_file_path}"'
+                    subprocess.Popen(cmd, shell=True)
+                    launched = True
+                    break
+            if not launched:
+                subprocess.Popen(
+                    [python, script, launch_arg, "--config-file", config_file_path],
+                    stdout=subprocess.DEVNULL,
+                    stderr=subprocess.DEVNULL,
+                    stdin=subprocess.DEVNULL,
+                    start_new_session=True,
+                )
+        elif system == "Darwin":
+            cmd = f'open -a Terminal.app {python} {script} {launch_arg} --config-file "{config_file_path}"'
+            subprocess.Popen(cmd, shell=True)
+        else:
+            subprocess.Popen(
+                [python, script, launch_arg, "--config-file", config_file_path],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                stdin=subprocess.DEVNULL,
+                start_new_session=True,
+            )
+    except Exception as e:
+        print(f"Failed to launch instance: {e}")
+        traceback.print_exc()
+        try:
+            os.unlink(config_file_path)
+        except:
+            pass
+
+
+def interactive_collect():
+    servers = []
+    clients = []
+    while True:
+        print("\n--- Add New Instance ---")
+        while True:
+            type_choice = input("Select type (0=Server, 1=Client): ").strip()
+            if type_choice in ("0", "1"):
+                break
+            print("Invalid input, please enter 0 or 1")
+        is_server = type_choice == "0"
+        while True:
+            setup_addr = input("Enter bind address and port (format host:port): ").strip()
+            try:
+                host, port = parse_addr_port(setup_addr)
+                break
+            except:
+                print("Invalid format, please retry")
+        if is_server:
+            config = {"host": host, "port": port}
+            servers = [config]
+            print(f"Server config set to: {host}:{port}")
+        else:
+            while True:
+                conn_addr = input(
+                    "Enter server address and port to connect (format host:port): "
+                ).strip()
+                try:
+                    srv_host, srv_port = parse_addr_port(conn_addr)
+                    break
+                except:
+                    print("Invalid format, please retry")
+            config = {"client_host": host, "client_port": port, "host": srv_host, "port": srv_port}
+            clients = [config]
+            print(f"Client config set to: local {host}:{port} -> server {srv_host}:{srv_port}")
+        cont = input("Continue adding more instances? (Y/N): ").strip().lower()
+        if cont != "y":
+            break
+    return servers, clients
+
+
+def generate_configs_from_args(args):
+    if args.setup_num > 1:
+        print("Warning: --setup_num is ignored because only one instance per type is allowed.")
+    host, port = parse_addr_port(args.setup_addr_port)
+    if args.type == 0:
+        config = {"host": host, "port": port}
+        return [config], []
+    else:
+        srv_host, srv_port = parse_addr_port(args.connect_addr_port)
+        config = {"client_host": host, "client_port": port, "host": srv_host, "port": srv_port}
+        return [], [config]
+
+
+def run_launched_instance(instance_type, config_file_path):
+    try:
+        with open(config_file_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+        try:
+            os.unlink(config_file_path)
+        except:
+            pass
+        if instance_type == "server":
+            server = TCP_Server_Base(**config)
+        else:
+            client = TCP_Client_Base(**config)
+    except Exception as e:
+        print(f"Failed to start instance: {e}")
+        traceback.print_exc()
+        input("Press any key to exit...")
+        sys.exit(1)
+
+
+def main():
+    if "--launch_server" in sys.argv:
+        idx = sys.argv.index("--launch_server")
+        try:
+            cfg_idx = sys.argv.index("--config-file", idx)
+            if cfg_idx + 1 < len(sys.argv):
+                config_file = sys.argv[cfg_idx + 1]
+                run_launched_instance("server", config_file)
+            else:
+                print("Error: missing --config-file argument")
+                sys.exit(1)
+        except ValueError:
+            print("Error: missing --config-file argument")
+            sys.exit(1)
+        return
+    if "--launch_client" in sys.argv:
+        idx = sys.argv.index("--launch_client")
+        try:
+            cfg_idx = sys.argv.index("--config-file", idx)
+            if cfg_idx + 1 < len(sys.argv):
+                config_file = sys.argv[cfg_idx + 1]
+                run_launched_instance("client", config_file)
+            else:
+                print("Error: missing --config-file argument")
+                sys.exit(1)
+        except ValueError:
+            print("Error: missing --config-file argument")
+            sys.exit(1)
+        return
+    parser = argparse.ArgumentParser(description="Flow Setup Launcher")
+    parser.add_argument("--type", type=int, choices=[0, 1], help="0=Server, 1=Client")
+    parser.add_argument("--setup_addr_port", type=str, help="Bind address and port (host:port)")
+    parser.add_argument(
+        "--connect_addr_port", type=str, help="Server address and port to connect (client required)"
+    )
+    parser.add_argument(
+        "--setup_num", type=int, default=1, help="Number of instances to launch (only 1 is allowed)"
+    )
+    args = parser.parse_args()
+    if args.type is not None:
+        if args.type == 0 and args.connect_addr_port is not None:
+            print("Error: --connect_addr_port cannot be used in Server mode")
+            sys.exit(1)
+        if args.type == 1 and (args.setup_addr_port is None or args.connect_addr_port is None):
+            print("Error: Client mode requires both --setup_addr_port and --connect_addr_port")
+            sys.exit(1)
+        if args.type == 0 and args.setup_addr_port is None:
+            print("Error: Server mode requires --setup_addr_port")
+            sys.exit(1)
+        servers, clients = generate_configs_from_args(args)
+        save_config(servers, clients)
+        for cfg in servers:
+            launch_instance(cfg, "server")
+        for cfg in clients:
+            launch_instance(cfg, "client")
+        return
+    if os.path.exists("setup.json"):
+        choice = input("setup.json exists. Overwrite configuration data? (Y/N): ").strip().lower()
+        if choice == "n":
+            config_data = load_existing_config()
+            for cfg in config_data.get("servers", []):
+                launch_instance(cfg, "server")
+            for cfg in config_data.get("clients", []):
+                launch_instance(cfg, "client")
+            return
+    else:
+        choice = "y"
+    servers, clients = interactive_collect()
+    save_config(servers, clients)
+    for cfg in servers:
+        launch_instance(cfg, "server")
+    for cfg in clients:
+        launch_instance(cfg, "client")
+
+
+if __name__ == "__main__":
+    main()