restrict.py

#!/usr/bin/env python3

import csv
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed

import paramiko

STUDENTS_FILE = "students.csv"
TARGETS_FILE = "eval_targets.txt"
NEW_PASSWORD = "restricted"


def change_password(student):
    name = student["name"]
    urn = student["urn"]
    ip = student["ip"]
    username = student["username"]
    password = student["password"]

    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(
            hostname=ip,
            username=username,
            password=password,
            timeout=15,
            allow_agent=False,
            look_for_keys=False,
        )

        _, stdout, _ = client.exec_command(
            f"echo '{username}:{NEW_PASSWORD}' | sudo chpasswd && echo OK || echo FAIL"
        )
        result = stdout.read().decode().strip()
        client.close()

        print(f"  {'Done ' if result == 'OK' else 'FAIL '} {urn} ({name}) at {ip}")

    except Exception as e:
        print(f"  ERROR {urn} ({name}) at {ip} → {e}")


def load_students(targets=None):
    students = []
    with open(STUDENTS_FILE) as f:
        for row in csv.DictReader(f):
            if targets is None or row["urn"] in targets:
                students.append(row)
    return students


def load_targets():
    if not os.path.exists(TARGETS_FILE):
        return None
    with open(TARGETS_FILE) as f:
        return {line.strip() for line in f if line.strip()}


def main():
    run_all = "--all" in sys.argv

    if run_all:
        students = load_students()
        print(f"Restricting ALL {len(students)} instance(s)...")
    else:
        targets = load_targets()
        if not targets:
            print(
                f"Error: {TARGETS_FILE} is empty. Use --all to restrict all instances."
            )
            sys.exit(1)
        students = load_students(targets)
        print(f"Restricting {len(students)} targeted instance(s)...")

    if not students:
        print("No matching students found.")
        sys.exit(1)

    print()
    with ThreadPoolExecutor(max_workers=len(students)) as executor:
        futures = {executor.submit(change_password, s): s for s in students}
        for future in as_completed(futures):
            future.result()

    print("\nDone. Password changed to 'restricted' on all targeted instances.")


if __name__ == "__main__":
    main()