change cluster.py

stevensusas · Nov 19, 2024 · b98e558 · b98e558
1 parent 933467d
commit b98e558
Showing 1 changed file with 63 additions and 133 deletions.
diff --git a/Control/cluster.py b/Control/cluster.py
@@ -8,7 +8,7 @@
 CORS(app)
 
 
-@app.route("/start-machine", methods=["POST"])
+@app.route('/start-machine', methods=['POST'])
 def start_machine():
     try:
         data = request.get_json()
@@ -19,140 +19,70 @@ def start_machine():
         if not vm_name:
             return jsonify({"status": "error", "message": "VM name is required"}), 400
 
-        start_script = f"""
-        tell application "UTM"
-            set vm to virtual machine named "{vm_name}"
-            start vm
-        end tell
-        """
-
-        print(f"Executing AppleScript: {start_script}")
-
-        process = subprocess.Popen(
-            ["osascript", "-e", start_script],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-        )
-        output, error = process.communicate()
-
-        print(f"Process return code: {process.returncode}")
-        print(f"Process output: {output}")
-        print(f"Process error: {error}")
-
-        if process.returncode != 0 or error:
-            error_message = error or "Unknown error occurred"
-            print(f"Error starting VM: {error_message}")
-            return (
-                jsonify(
-                    {
-                        "status": "error",
-                        "message": error_message,
-                        "details": {
-                            "returncode": process.returncode,
-                            "output": output,
-                            "error": error,
-                        },
-                    }
-                ),
-                500,
-            )
-
-        return jsonify(
-            {
-                "status": "success",
-                "message": "VM started successfully",
-                "details": {"output": output},
-            }
-        )
-
-    except Exception as e:
-        print(f"Exception in start_machine:")
-        print(traceback.format_exc())
-        return (
-            jsonify(
-                {
-                    "status": "error",
-                    "message": str(e),
-                    "traceback": traceback.format_exc(),
-                }
-            ),
-            500,
-        )
-
-
-@app.route("/stop-machine", methods=["POST"])
+    start_script = f'''
+    tell application "UTM"
+        set vm to virtual machine named "{vm_name}"
+        start vm    
+    end tell
+    '''
+
+    process = subprocess.Popen(['osascript', '-e', start_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    output, error = process.communicate()
+
+    if error:
+        return jsonify({"status": "error", "message": error.decode()}), 500
+    else:
+        return jsonify({"status": "success", "message": "VM started successfully"})
+
+@app.route('/stop-machine', methods=['POST'])
 def stop_machine():
-    try:
-        data = request.get_json()
-        vm_name = data.get("vm_name")
-
-        print(f"Attempting to stop VM: {vm_name}")
+    data = request.get_json()
+    vm_name = data.get("vm_name")
+    stop_type = data.get("stop_type", "suspend with saving")  # Default to suspend with saving
 
         if not vm_name:
             return jsonify({"status": "error", "message": "VM name is required"}), 400
 
-        stop_script = f"""
-        tell application "UTM"
-            set vm to virtual machine named "{vm_name}"
-            suspend vm
-        end tell
-        """
-
-        print(f"Executing AppleScript: {stop_script}")
-
-        process = subprocess.Popen(
-            ["osascript", "-e", stop_script],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-        )
-        output, error = process.communicate()
-
-        print(f"Process return code: {process.returncode}")
-        print(f"Process output: {output}")
-        print(f"Process error: {error}")
-
-        if process.returncode != 0 or error:
-            error_message = error or "Unknown error occurred"
-            print(f"Error stopping VM: {error_message}")
-            return (
-                jsonify(
-                    {
-                        "status": "error",
-                        "message": error_message,
-                        "details": {
-                            "returncode": process.returncode,
-                            "output": output,
-                            "error": error,
-                        },
-                    }
-                ),
-                500,
-            )
-
-        return jsonify(
-            {
-                "status": "success",
-                "message": "VM stopped successfully",
-                "details": {"output": output},
-            }
-        )
-
-    except Exception as e:
-        print(f"Exception in stop_machine:")
-        print(traceback.format_exc())
-        return (
-            jsonify(
-                {
-                    "status": "error",
-                    "message": str(e),
-                    "traceback": traceback.format_exc(),
-                }
-            ),
-            500,
-        )
-
-
-if __name__ == "__main__":
-    app.run(port=5000, debug=True)
+    # Construct the AppleScript command based on stop_type
+    stop_script = f'''
+    tell application "UTM"
+        set vm to virtual machine named "{vm_name}"
+        {stop_type} vm
+    end tell
+    '''
+
+    process = subprocess.Popen(['osascript', '-e', stop_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    output, error = process.communicate()
+
+    if error:
+        return jsonify({"status": "error", "message": error.decode()}), 500
+    else:
+        return jsonify({"status": "success", "message": f"VM paused and saved with {stop_type} successfully"})
+
+@app.route('/reset-machine', methods=['POST'])
+def reset_machine():
+    data = request.get_json()
+    vm_name = data.get("vm_name")
+
+    if not vm_name:
+        return jsonify({"status": "error", "message": "VM name is required"}), 400
+
+    # Construct the AppleScript command to stop the VM
+    reset_script = f'''
+    tell application "UTM"
+        set vm to virtual machine named "{vm_name}"
+        stop vm
+    end tell
+    '''
+
+    process = subprocess.Popen(['osascript', '-e', reset_script], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    output, error = process.communicate()
+
+    if error:
+        return jsonify({"status": "error", "message": error.decode()}), 500
+    else:
+        return jsonify({"status": "success", "message": "VM stopped successfully"})
+
+
+if __name__ == '__main__':
+    app.run(port=5000)