import subprocess
import os

class SSHClient:
    def __init__(self, host, user="root", key_path=None, password=None):
        self.host = host
        self.user = user
        self.key_path = os.path.expanduser(key_path) if key_path else None
        self.password = password

    def run(self, cmd, capture=True):
        ssh_cmd = ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null"]
        
        if self.key_path:
            ssh_cmd += ["-i", self.key_path]
        
        target = f"{self.user}@{self.host}"
        
        if self.password:
            full_cmd = ["sshpass", "-p", self.password] + ssh_cmd + [target, cmd]
        else:
            full_cmd = ssh_cmd + [target, cmd]

        result = subprocess.run(
            full_cmd,
            capture_output=capture,
            text=True
        )
        return result

    def scp_to(self, local_path, remote_path):
        scp_cmd = ["scp", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null"]
        if self.key_path:
            scp_cmd += ["-i", self.key_path]
        
        target = f"{self.user}@{self.host}:{remote_path}"
        
        if self.password:
            full_cmd = ["sshpass", "-p", self.password] + scp_cmd + [local_path, target]
        else:
            full_cmd = scp_cmd + [local_path, target]

        return subprocess.run(full_cmd, capture_output=True)