Update app.py

app.py

@@ -1,64 +1,167 @@
-from flask import Flask, request, jsonify
 import os
-import subprocess
 import tempfile
-import json
+import shutil
+from zipfile import ZipFile
+import logging
+import psutil
+import subprocess
+from flask import Flask, request, jsonify, render_template, send_file
 import multiprocessing
 
-app = Flask(__name__)
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Dictionary to store donated CPU information
-donated_cpus = {}
-total_donated_cpus = 0
+# Initialize Flask app
+app = Flask(__name__)
 
-@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}"})
+connected_cpus = {"localhost": {"cpu_count": psutil.cpu_count(logical=False), "usage": 0.0}}
 
-@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
+# 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('/execute_script', methods=['POST'])
-def execute_script():
+# Endpoint to handle file uploads and script execution
+@app.route('/upload', methods=['POST'])
+def handle_upload():
     try:
-        data = request.get_json()
-        script_content = data['script_content']
+        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
         
-        # Save the script to a temporary file
-        script_file = tempfile.NamedTemporaryFile(delete=False, suffix='.py')
-        script_file.write(script_content.encode())
-        script_file.close()
+        files = request.files.getlist('file')
+        script_content = request.form['script_content']
         
-        # 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)
+        # Create a temporary directory to store uploaded files
+        temp_dir = tempfile.mkdtemp()
         
-        os.remove(script_file.name)
-        return jsonify({
-            "stdout": result.stdout,
-            "stderr": result.stderr,
-            "returncode": result.returncode
-        })
+        # 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)}"})
+    
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        logger.error(f"Error in handle_upload: {e}")
+        return jsonify({"status": "error", "message": 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 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})
+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
+
+# Endpoint to donate CPU resources
+@app.route('/donate_cpu', methods=['POST'])
+def donate_cpu():
+    try:
+        data = request.json
+        host = data['host']
+        cpu_count = data['cpu_count']
+        connected_cpus[host] = {"cpu_count": cpu_count, "usage": 0.0}
+        return jsonify({"status": "success", "message": f"CPU resources from {host} donated successfully."})
+    except Exception as e:
+        logger.error(f"Error in donate_cpu: {e}")
+        return jsonify({"status": "error", "message": str(e)}), 500
+
+# Endpoint to update CPU usage
+@app.route('/update_cpu_usage', methods=['POST'])
+def update_cpu_usage():
+    try:
+        data = request.json
+        host = data['host']
+        usage = data['usage']
+        if host in connected_cpus:
+            connected_cpus[host]['usage'] = usage
+            return jsonify({"status": "success", "message": f"CPU usage from {host} updated successfully."})
+        else:
+            return jsonify({"status": "error", "message": f"Host {host} not found."}), 404
+    except Exception as e:
+        logger.error(f"Error in update_cpu_usage: {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)
 
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860)
+    app.run(host='0.0.0.0', port=7860, threaded=True)