def test_host_command(self):
configs = SSHConfig(sample)
self.assertEqual("ssh 203.0.113.76", configs.get("server1").command())
self.assertEqual("ssh -p 2202 203.0.113.76",
configs.get("server_cmd_1").command())
self.assertEqual("ssh [email protected]",
configs.get("server_cmd_2").command())
self.assertEqual("ssh -p 2202 [email protected]",
configs.get("server_cmd_3").command())
def test_write(self):
configs = SSHConfig(sample)
configs.add(new_host)
new_sample_path = os.path.join(os.path.dirname(__file__), "sample_new")
configs.write(filename=new_sample_path)
new_config = SSHConfig(new_sample_path)
os.remove(new_sample_path)
self.assertEqual("server2", new_config.get("server2").name)
def test_update(self):
configs = SSHConfig(sample)
configs.update("server1", {"IdentityFile": "~/.ssh/id_rsa_new"})
self.assertRaises(AttributeError, configs.update, "server1", [])
self.assertEqual(
configs.get("server1").IdentityFile, "~/.ssh/id_rsa_new")
attrs = {
"HostName": "example.com",
"User": "******",
"Port": 22,
"IdentityFile": "~/.ssh/id_rsa",
"ServerAliveInterval": 10,
}
configs.update("server1", attrs)
for key, value in attrs.items():
self.assertEqual(getattr(configs.get("server1"), key), value)
def prepare_ssh_config(cluster_id, profile, public_dns):
"""Add/edit the ssh configuration belonging to the given cluster in ~/.ssh/config
Args:
cluster_id (str): Cluster ID
profile (str): Databricks CLI profile string
public_dns (str): Public DNS/IP address
"""
backup_path = "~/.databrickslabs_jupyterlab/ssh_config_backup"
config = os.path.expanduser("~/.ssh/config")
if not os.path.exists(os.path.expanduser(backup_path)):
os.makedirs(os.path.expanduser(backup_path))
print_warning(" => ~/.ssh/config will be changed")
backup = "%s/config.%s" % (backup_path, time.strftime("%Y-%m-%d_%H-%M-%S"))
print_warning(
" => A backup of the current ~/.ssh/config has been created")
print_warning(" => at %s" % backup)
shutil.copy(config, os.path.expanduser(backup))
try:
sc = SSHConfig.load(config)
except:
sc = SSHConfig(config)
hosts = [h.name for h in sc.hosts()]
if cluster_id in hosts:
host = sc.get(cluster_id)
host.set("HostName", public_dns)
host.set("ServerAliveInterval", 30)
host.set("ServerAliveCountMax", 5760)
host.set("ConnectTimeout", 5)
print(" => Added ssh config entry or modified IP address:\n")
print(textwrap.indent(str(host), " "))
else:
# ServerAliveInterval * ServerAliveCountMax = 48h
attrs = {
"HostName": public_dns,
"IdentityFile": "~/.ssh/id_%s" % profile,
"Port": 2200,
"User": "******",
"ServerAliveInterval": 30,
"ServerAliveCountMax": 5760,
"ConnectTimeout": 5,
}
host = Host(name=cluster_id, attrs=attrs)
print(" => Adding ssh config to ~/.ssh/config:\n")
print(textwrap.indent(str(host), " "))
sc.append(host)
sc.write()
add_known_host(public_dns)
def test_remove(self):
config = SSHConfig(sample)
config.remove("server1")
with self.assertRaises(NameError):
config.get("server1")
def test_get_host(self):
configs = SSHConfig(sample)
self.assertEqual("server1", configs.get("server1").name)
with self.assertRaises(NameError):
configs.get("NoExist")