123 lines
3.9 KiB
Python
Executable File
123 lines
3.9 KiB
Python
Executable File
#!/usr/bin/python3
|
|
# ==============================================================#
|
|
# File : ssh
|
|
# Desc : setup vagrant ssh config
|
|
# Ctime : 2021-04-20
|
|
# Mtime : 2025-09-11
|
|
# Path : vagrant/ssh
|
|
# License : Apache-2.0 @ https://pigsty.io/docs/about/license/
|
|
# Copyright : 2018-2026 Ruohang Feng / Vonng (rh@vonng.com)
|
|
# ==============================================================#
|
|
# this script will generate vagrant ssh config to ~/.ssh
|
|
# which allows you to ssh to VMs via both name and IP addr
|
|
|
|
import os, re, sys, subprocess
|
|
|
|
|
|
# get the node name and ip mapping from Vagrantfile
|
|
def parse_vagrant_spec(vagrant_filepath):
|
|
specs = []
|
|
result = {}
|
|
with open(vagrant_filepath, 'r') as f:
|
|
raw_lines = f.readlines()
|
|
|
|
activate = False
|
|
for line in raw_lines:
|
|
if line.startswith('Specs'):
|
|
activate = True
|
|
continue
|
|
elif line.startswith(']'):
|
|
activate = False
|
|
break
|
|
if activate and '"name"' in line and '"ip"' in line:
|
|
name = re.findall(r'"name"\s*=>\s*"([^"]+)"', line)[0]
|
|
ip = re.findall(r'"ip"\s*=>\s*"([^"]+)"', line)[0]
|
|
result[name] = ip
|
|
specs.append((ip, name))
|
|
return result, specs
|
|
|
|
|
|
# run vagrant ssh-config and gather result
|
|
def get_vagrant_ssh_config(vagrant_dir):
|
|
print("cd %s && vagrant ssh-config" % vagrant_dir)
|
|
print("this may take several seconds....")
|
|
os.chdir(vagrant_dir)
|
|
command = 'vagrant ssh-config'
|
|
config = subprocess.check_output(command, shell=True, universal_newlines=True)
|
|
return config
|
|
|
|
|
|
def generate_ssh_config(config, mapping):
|
|
extra_config = config
|
|
for name, ip in mapping.items():
|
|
src = 'Host %s' % name
|
|
dst = 'Host %s' % ip
|
|
extra_config = extra_config.replace(src, dst)
|
|
final_config = '# ssh access via nodename\n\n\n' + config + '\n\n\n' + '# SSH Access via IP address\n\n' + extra_config
|
|
return final_config
|
|
|
|
|
|
def write_ssh_config(config, name="pigsty"):
|
|
config_name = '%s_config' % name
|
|
include_cmd = 'Include ~/.ssh/%s_config' % name
|
|
main_config_path = os.path.join(os.environ['HOME'], '.ssh', 'config')
|
|
extra_config_path = os.path.join(os.environ['HOME'], '.ssh', config_name)
|
|
|
|
# write extra config
|
|
print("write extra ssh config [%s] to %s" % (name, extra_config_path))
|
|
with open(extra_config_path, 'w') as f:
|
|
f.write(config)
|
|
os.chmod(extra_config_path, 0o600)
|
|
|
|
with open(main_config_path, 'r') as f:
|
|
ssh_config = f.read()
|
|
if include_cmd in ssh_config:
|
|
print("include cmd already exists in ~/.ssh/config")
|
|
return
|
|
|
|
print("write include %s command to ~/.ssh/config" % name)
|
|
with open(main_config_path, 'a') as f:
|
|
f.write('\n\n# %s\n%s\n' % (config_name, include_cmd))
|
|
os.chmod(main_config_path, 0o600)
|
|
|
|
|
|
def get_config_name():
|
|
config_name = 'pigsty'
|
|
if len(sys.argv) > 1:
|
|
config_name = sys.argv[1]
|
|
if not re.match('^[a-z0-9]+$', config_name):
|
|
print("config name must match regexp: [a-z0-9]+")
|
|
sys.exit(1)
|
|
return config_name
|
|
|
|
|
|
def main():
|
|
# generate config name, pigsty by default
|
|
config_name = get_config_name()
|
|
|
|
# check vagrantfile exists
|
|
vagrant_dir = os.path.dirname(os.path.realpath(__file__))
|
|
os.chdir(vagrant_dir)
|
|
vagrant_file = os.path.join(vagrant_dir, 'Vagrantfile')
|
|
if not os.path.exists(vagrant_file):
|
|
print("Vagrantfile not found in %s" % vagrant_file)
|
|
sys.exit(1)
|
|
|
|
# parse vagrantfile and generate ssh config
|
|
mapping, specs = parse_vagrant_spec(vagrant_file)
|
|
print("\nVagrant nodes:\n")
|
|
for ip, name in specs:
|
|
print("%-16s %s" % (ip, name))
|
|
print("\n\n")
|
|
config = get_vagrant_ssh_config(vagrant_dir)
|
|
config = generate_ssh_config(config, mapping)
|
|
print(config)
|
|
|
|
# write ssh config to ~/.ssh/
|
|
write_ssh_config(config, config_name)
|
|
|
|
|
|
# main
|
|
if __name__ == '__main__':
|
|
main()
|