reportmonster/classes/connection.py

50 lines
1.6 KiB
Python

import paramiko as pikuniku # :P
from paramiko.client import SSHClient, WarningPolicy
from sshtunnel import SSHTunnelForwarder
from typing import Union, Optional
from contextlib import closing
import socket
class Connection:
"""Class representing an SSH/SFTP connection to a Vessel
"""
def __init__(self, vessel):
"""Initialize a new Connection to a Vessel
Args:
vessel (classes.vessel.Vessel): Vessel object to open connection to
"""
self._vessel = vessel
self._client = SSHClient()
self._client.load_system_host_keys()
self._client.set_missing_host_key_policy(WarningPolicy)
self._client.connect(vessel.host, 22, vessel.ssh_username,
vessel.ssh_password, timeout=vessel.ssh_timeout,
passphrase=vessel.ssh_passphrase)
self._transport = self._client.get_transport()
self._transport.set_keepalive(10)
self._sftp = self._client.open_sftp()
self._process = None
def forward_tcp(self, remote=3306):
self._process = SSHTunnelForwarder(
(self._vessel.host, 22),
ssh_username=self._vessel.ssh_username,
ssh_private_key_password=self._vessel.ssh_passphrase,
remote_bind_address=("127.0.0.1", remote))
self._process.start()
return self._process.local_bind_port
def __del__(self):
"""Close SSH connection when ending Connection
"""
self._client.close()
if self._process:
self._process.close()