Skip to content

Commit

Permalink
Merge pull request #1 from zubairahm3d/feature/shahzaib
Browse files Browse the repository at this point in the history 
Feature/shahzaib
  • Loading branch information
@HypertextAssassin0273
HypertextAssassin0273 authored Nov 26, 2024
2 parents 42d60f2 + 832a9e2 commit 53d3a0b
Show file tree
Hide file tree
Showing 2 changed files with 100 additions and 21 deletions.
69 changes: 48 additions & 21 deletions app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
def cleanup_on_exit():
"""calls the cleanup route at app shutdown to remove any spawned containers."""
with app.test_client() as flask_client:
response = flask_client.post('/cleanup') # call the /cleanup route
response = flask_client.post('/remove_all_containers') # calls remove_all_containers route
app.logger.info(f"Cleanup result: {response.get_data(as_text=True)}")

atexit.register(cleanup_on_exit) # register 'cleanup' to run when the app quits or receives SIGINT
Expand Down Expand Up @@ -207,28 +207,31 @@ def run_ansible(hosts, command):
finally:
cleanup_files([inventory_file_path, playbook_file_path])

def remove_containers(containers, type="spawned"):
def remove_or_stop_containers(containers, action):
"""
Stops and removes specified containers.
Stops or removes the specified containers based on the action parameter.
"""
for container in containers:
try:
# if container.status != 'exited' and container.status != 'created':
container.stop() # [NOTE] default timeout is culprit for unssuccessful removal in 1st attempt
container.remove()
app.logger.info(f"{type} container {container.id[:12]} removed.")
if container.status == 'running':
container.stop() # [NOTE] default timeout is culprit for unssuccessful removal in 1st attempt
if action == 'remove':
container.remove()
app.logger.info(f"spawned container {container.id[:12]} {'removed' if action == 'remove' else 'stopped'}")

except Exception as e:
app.logger.warning(f"Failed to remove {type} container {container.id[:12]}: {str(e)}")
app.logger.warning(f"Failed to {action} spawned container {container.id[:12]}: {str(e)}")


## ROUTES FOR FLASK APP ##
# HTTP status codes: 200 - OK (default), 400 - Bad Request, 404 - Not Found, 500 - Internal Server Error

@app.route('/')
def index():
def MAIN_INDEX_ROUTE():
return render_template('index.html')

@app.route('/spawn', methods=['POST'])
def spawn_machines():
def SPAWN_MACHINES_ROUTE():
try:
num_machines = request.form.get('num_machines', type=int)
if not isinstance(num_machines, int) or num_machines <= 0:
Expand All @@ -238,7 +241,7 @@ def spawn_machines():
containers = [] # Track spawned containers
machine_info = [] # Track machine details for display

for i in range(num_machines):
for _ in range(num_machines):
try:
container = spawn_container(public_key)
host_port = wait_for_container(container)
Expand Down Expand Up @@ -268,7 +271,7 @@ def spawn_machines():


@app.route('/run_command', methods=['POST'])
def run_command():
def RUN_COMMAND_ROUTE():
try:
command = request.form['command']
if not command or ';' in command or '&&' in command: # prevent command injection
Expand Down Expand Up @@ -305,21 +308,45 @@ def run_command():
app.logger.error(f"Error running command: {str(e)}")
return jsonify({"error": str(e)}), 500

@app.route('/cleanup', methods=['POST'])
def cleanup_containers():
"""
Cleans up all spawned containers and their associated data.
"""
@app.route('/<action>_all_containers', methods=['POST'])
def STOP_OR_REMOVE_ALL_CONTAINERS_ROUTE(action):
try:
# get all spawned containers ([NOTE] including orphaned?)
if not session.get('machine_info'): # safeguard against empty session data
return jsonify({"status": "No machines spawned. Nothing to stop or remove."})

if action not in ['stop', 'remove']: # validate the action parameter
return jsonify({"error": f"Invalid action: {action}. Use 'stop' or 'remove'."}), 400

containers = client.containers.list(all=True, filters={"label": "flask_app=spawned_container"})
remove_containers(containers)
remove_or_stop_containers(containers, action)

session.pop('machine_info', None) # clear session data
return jsonify({"status": "Cleanup completed successfully."})
return jsonify({"status": f"All containers {'stopped and' if action == 'stop' else ''} removed successfully."})

except Exception as e:
app.logger.error(f"Error {'stopping' if action == 'stop' else 'removing' } all containers: {str(e)}")
return jsonify({"error": str(e)}), 500


@app.route('/<action>_container', methods=['POST'])
def STOP_OR_REMOVE_CONTAINER_ROUTE():
try:
container_id = request.form.get('container_id')
if not container_id:
return jsonify({"error": "No container ID provided."}), 400

container = client.containers.get(container_id)
remove_or_stop_containers([container], 'stop')

# Remove the stopped container from the session data
session['machine_info'] = [machine for machine in session.get('machine_info', []) if machine['container_id'] != container_id]

return jsonify({"status": f"Container {container_id[:12]} stopped and removed successfully."})

except docker.errors.NotFound:
return jsonify({"error": f"Container {container_id[:12]} not found."}), 404
except Exception as e:
app.logger.error(f"Error during cleanup: {str(e)}")
app.logger.error(f"Error stopping container {container_id[:12]}: {str(e)}")
return jsonify({"error": str(e)}), 500


Expand Down
52 changes: 52 additions & 0 deletions templates/index.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
colors: {
primary: '#3B82F6',
secondary: '#10B981',
danger: '#EF4444',
}
}
}
Expand All @@ -35,6 +36,9 @@ <h1 class="text-4xl font-bold text-center mb-8 text-gray-800">Ansible Manager</h
<button onclick="spawnMachines()" class="bg-primary text-white font-semibold px-6 py-2 rounded-md hover:bg-blue-600 transition duration-300 ease-in-out transform hover:scale-105">
Spawn Machines
</button>
<button onclick="stopAllMachines()" class="bg-danger text-white font-semibold px-6 py-2 rounded-md hover:bg-red-600 transition duration-300 ease-in-out transform hover:scale-105">
Stop All Machines
</button>
</div>
</div>
<div id="machineInfo" class="space-y-4"></div>
Expand Down Expand Up @@ -115,6 +119,11 @@ <h3 class="text-2xl font-semibold mb-4 text-gray-800">Machine ${index + 1}</h3>
<p class="text-gray-600"><span class="font-medium text-gray-700">SSH Command:</span> <code class="bg-gray-100 px-2 py-1 rounded">${machine.host_command}</code></p>
<p class="text-gray-600"><span class="font-medium text-gray-700">SSH Status:</span> <span class="${machine.ssh_status === 'Ready' ? 'text-green-600' : 'text-red-600'} font-semibold">${machine.ssh_status}</span></p>
</div>
<div class="mt-4">
<button onclick="stopMachine('${machine.container_id}')" class="bg-danger text-white font-semibold px-4 py-2 rounded-md hover:bg-red-600 transition duration-300 ease-in-out transform hover:scale-105">
Stop Machine
</button>
</div>
`;
machineInfo.appendChild(machineDiv);
});
Expand Down Expand Up @@ -160,6 +169,49 @@ <h3 class="text-2xl font-semibold mb-4 text-gray-800">Machine ${index + 1} (${ma
commandResults.appendChild(resultDiv);
});
}

function stopAllMachines() {
const machineInfo = document.getElementById('machineInfo');
machineInfo.innerHTML = '<div class="text-center text-xl font-semibold text-gray-600">Stopping all machines...</div>';

axios.post('/remove_all_containers') // [IMPROVEMENT] change to stop_all_machines for later
.then(function (response) {
if (response.data.error) {
machineInfo.innerHTML = `<div class="text-center text-xl font-semibold text-red-600">Error: ${response.data.error}</div>`;
return;
}
machineInfo.innerHTML = `<div class="text-center text-xl font-semibold text-green-600">${response.data.status}</div>`;
machineInfoData = [];
})
.catch(function (error) {
console.error('Error:', error);
machineInfo.innerHTML = `<div class="text-center text-xl font-semibold text-red-600">Error stopping machines: ${error.response?.data?.error || error.message}</div>`;
});
}

function stopMachine(containerId) {
const machineInfo = document.getElementById('machineInfo');
machineInfo.innerHTML = `<div class="text-center text-xl font-semibold text-gray-600">Stopping machine ${containerId}...</div>`;

axios.post('/remove_container', `container_id=${containerId}`, { // [IMPROVEMENT] change to stop_container for later
headers: {
'Content-Type': 'application/x-www-form-urlencoded'
}
})
.then(function (response) {
if (response.data.error) {
machineInfo.innerHTML = `<div class="text-center text-xl font-semibold text-red-600">Error: ${response.data.error}</div>`;
return;
}
machineInfoData = machineInfoData.filter(machine => machine.container_id !== containerId);
displayMachineInfo(machineInfoData);
machineInfo.insertAdjacentHTML('afterbegin', `<div class="text-center text-xl font-semibold text-green-600 mb-4">${response.data.status}</div>`);
})
.catch(function (error) {
console.error('Error:', error);
machineInfo.innerHTML = `<div class="text-center text-xl font-semibold text-red-600">Error stopping machine: ${error.response?.data?.error || error.message}</div>`;
});
}
</script>
</body>
</html>
Expand Down

0 comments on commit 53d3a0b

Please sign in to comment.