Update app.py

app.py

@@ -1,137 +1,64 @@
+from flask import Flask, request, jsonify
 import os
-import tempfile
-import shutil
-from zipfile import ZipFile
-import logging
-import psutil
 import subprocess
-from flask import Flask, request, jsonify, render_template, send_file
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+import tempfile
+import json
+import multiprocessing
 
-# Initialize Flask app
 app = Flask(__name__)
 
-connected_cpus = {"localhost": {"cpu_count": psutil.cpu_count(logical=False), "usage": 0.0}}
+# Dictionary to store donated CPU information
+donated_cpus = {}
+total_donated_cpus = 0
 
-# Define the target function for multiprocessing
-def target_function(script_path, folder_path):
-    output_log = tempfile.TemporaryFile(mode='w+t')
-    try:
-        result = subprocess.run(['python', script_path], cwd=folder_path, stdout=output_log, stderr=subprocess.STDOUT)
-        output_log.seek(0)
-        log_output = output_log.read()
-    except Exception as e:
-        log_output = str(e)
-    finally:
-        output_log.close()
-    return log_output
+@app.route('/donate_cpu', methods=['POST'])
+def donate_cpu():
+    global total_donated_cpus
+    data = request.get_json()
+    host = data['host']
+    cpu_count = data['cpu_count']
+    donated_cpus[host] = {"cpu_count": cpu_count, "usage": 0.0}
+    total_donated_cpus += cpu_count
+    return jsonify({"status": "success", "message": f"CPU donated by {host}"})
+
+@app.route('/update_cpu_usage', methods=['POST'])
+def update_cpu_usage():
+    data = request.get_json()
+    host = data['host']
+    usage = data['usage']
+    if host in donated_cpus:
+        donated_cpus[host]['usage'] = usage
+        return jsonify({"status": "success", "message": f"Updated CPU usage for {host}: {usage}%"})
+    return jsonify({"status": "error", "message": "Host not found"}), 404
 
-# Endpoint to handle file uploads and script execution
-@app.route('/upload', methods=['POST'])
-def handle_upload():
+@app.route('/execute_script', methods=['POST'])
+def execute_script():
     try:
-        if 'file' not in request.files or 'script_content' not in request.form:
-            return jsonify({"status": "error", "message": "File or script content not provided"}), 400
+        data = request.get_json()
+        script_content = data['script_content']
         
-        files = request.files.getlist('file')
-        script_content = request.form['script_content']
+        # Save the script to a temporary file
+        script_file = tempfile.NamedTemporaryFile(delete=False, suffix='.py')
+        script_file.write(script_content.encode())
+        script_file.close()
         
-        # Create a temporary directory to store uploaded files
-        temp_dir = tempfile.mkdtemp()
+        # Run the script using the donated CPUs
+        cpu_count = total_donated_cpus or 1  # Default to 1 if no CPUs are donated
+        result = subprocess.run(['python', script_file.name], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=60)
         
-        # Save the uploaded files to the temporary directory
-        folder_path = os.path.join(temp_dir, 'uploaded_folder')
-        os.makedirs(folder_path, exist_ok=True)
-        for file_obj in files:
-            file_path = os.path.join(folder_path, file_obj.filename)
-            file_obj.save(file_path)
-        
-        # Save the script content to a file
-        script_path = os.path.join(folder_path, 'user_script.py')
-        with open(script_path, 'w') as script_file:
-            script_file.write(script_content)
-        
-        # Run the script using multiprocessing
-        log_output = run_script(script_path, folder_path)
-        
-        # Create a zip file of the entire folder
-        zip_path = os.path.join(temp_dir, 'output_folder.zip')
-        with ZipFile(zip_path, 'w') as zipf:
-            for root, _, files in os.walk(folder_path):
-                for file in files:
-                    zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
-        
-        return jsonify({"status": "success", "log_output": log_output, "download_url": f"/download/{os.path.basename(zip_path)}"})
-    
+        os.remove(script_file.name)
+        return jsonify({
+            "stdout": result.stdout,
+            "stderr": result.stderr,
+            "returncode": result.returncode
+        })
     except Exception as e:
-        logger.error(f"Error in handle_upload: {e}")
-        return jsonify({"status": "error", "message": str(e)}), 500
+        return jsonify({"error": str(e)}), 500
 
-@app.route('/download/<filename>')
-def download_file(filename):
-    try:
-        return send_file(os.path.join(tempfile.gettempdir(), filename), as_attachment=True)
-    except Exception as e:
-        logger.error(f"Error in download_file: {e}")
-        return jsonify({"status": "error", "message": str(e)}), 500
-
-# Endpoint to get connected CPUs information
 @app.route('/cpu_info', methods=['GET'])
-def get_cpu_info():
-    try:
-        info = []
-        for host, data in connected_cpus.items():
-            info.append(f"{host}: {data['cpu_count']} CPUs, {data['usage']}% usage")
-        return jsonify({"status": "success", "cpu_info": "\n".join(info)})
-    except Exception as e:
-        logger.error(f"Error in get_cpu_info: {e}")
-        return jsonify({"status": "error", "message": str(e)}), 500
-
-# Endpoint to execute commands
-@app.route('/execute_command', methods=['POST'])
-def execute_command():
-    try:
-        command = request.form['command']
-        if not command:
-            return jsonify({"status": "error", "message": "No command provided"}), 400
-        
-        # Ensure commands are executed in a safe environment
-        allowed_commands = ['pip install']
-        if not any(command.startswith(cmd) for cmd in allowed_commands):
-            return jsonify({"status": "error", "message": "Command not allowed"}), 400
-        
-        output_log = tempfile.TemporaryFile(mode='w+t')
-        try:
-            result = subprocess.run(command.split(), stdout=output_log, stderr=subprocess.STDOUT)
-            output_log.seek(0)
-            log_output = output_log.read()
-        except Exception as e:
-            log_output = str(e)
-        finally:
-            output_log.close()
-        return jsonify({"status": "success", "log_output": log_output})
-    
-    except Exception as e:
-        logger.error(f"Error in execute_command: {e}")
-        return jsonify({"status": "error", "message": str(e)}), 500
-
-# Main interface
-@app.route('/')
-def index():
-    return render_template('index.html')
-
-def run_script(script_path, folder_path):
-    # Collect all available CPUs including the local host CPU
-    total_cpus = sum(cpu['cpu_count'] for cpu in connected_cpus.values()) + 1
-    
-    # Use multiprocessing to run the script
-    with multiprocessing.Pool(total_cpus) as pool:
-        log_outputs = pool.starmap(target_function, [(script_path, folder_path)] * total_cpus)
-
-    return '\n'.join(log_outputs)
+def cpu_info():
+    info = [{"host": host, "cpu_count": data['cpu_count'], "usage": data['usage']} for host, data in donated_cpus.items()]
+    return jsonify({"cpu_info": info})
 
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860, threaded=True)
+    app.run(host='0.0.0.0', port=7860)