import os
import json
import tornado.web
import subprocess



# New handler for your custom page
class AppHandler(tornado.web.RequestHandler):
    # Define the path to the templates directory
    HandlerPath = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),'WebSSH', 'templates') 
    template_folder = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'WebSSH', 'templates')

    def get(self):
        self.render(os.path.join(self.template_folder, 'index.html'))

    def post(self):
        username = self.get_argument("username")
        password = self.get_argument("password")

        # Check if the user already exists
        user_exists = subprocess.run(["id", "-u", username], capture_output=True)
        if user_exists.returncode == 0:
            # User exists, read the existing SSH key
            ssh_dir = f"/home/{username}/.ssh"
            private_key_path = f"{ssh_dir}/id_rsa"
            if os.path.exists(private_key_path):
                with open(private_key_path, "r") as file:
                    private_key = file.read()
            else:
                self.set_status(404)
                self.write("SSH key not found for existing user.")
                return
        else:
            """
            # Create the user directory and .ssh directory manually
            user_home = f"/home/{username}"
            ssh_dir = f"{user_home}/.ssh"
            os.makedirs(ssh_dir, exist_ok=True)

            # Generate SSH key pair for the new user
            subprocess.run(["ssh-keygen", "-t", "rsa", "-b", "2048", "-f", f"{ssh_dir}/id_rsa", "-N", ""])

            # Read the private key
            with open(f"{ssh_dir}/id_rsa", "r") as file:
                private_key = file.read()
            """
            self.set_status(404)
            self.write("SSH key not found for existing user.")
            return

        # Return the private key to the user
        self.set_header('Content-Type', 'application/octet-stream')
        self.set_header('Content-Disposition', f'attachment; filename=id_rsa')
        self.write(private_key)