from ..Playbooks.Playbook import FileCleanupTask, PlaybookCredentialCollection
from ...Logging.Logger import log, Severity
from ...Machine.MachineInterface import MachineInterface
from ..SetupStep import SetupStep
import socket

class TestConnection(SetupStep):
    def __init__(self, machine: MachineInterface, port: int):
        super().__init__()
        self.machine = machine
        self.port = port

    def run(self, file_cleanup_stack : list[FileCleanupTask], credentials: PlaybookCredentialCollection) -> bool:
        hostname = self.machine.get_hostname()
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(2)
            sock.connect((hostname, 22))
            sock.close()
            log(f"Connection test to {hostname} succeeded.")
            return True
        except (socket.gaierror, socket.timeout, ConnectionRefusedError, OSError) as e:
            log(f"Unable to reach {hostname}: {e}", Severity.ERROR)
            return False
        except Exception as e:
            log(f"Exception while trying to reach {hostname}: {e}", Severity.WTF)
            return False
