#!/usr/bin/env python
# coding: utf-8

# $Id: $
import socket
import sys
import os
import getpass

import paramiko
from storm import Storm

sconfig = Storm(os.path.expanduser("~/.ssh/config"))
sconfig.backup(os.path.expanduser("~/.ssh/config.bak"))
if len(sys.argv) < 2:
    username = raw_input("Enter username: ")
    host = raw_input("Enter host: ")
else:
    username, host = sys.argv[1].split('@')
name = host.split('.')[0]

if sconfig.is_host_in(name):
    print "Host %s already exists in .ssh/config, removing." % name
    sconfig.delete_entry(name)
print "Creating new entry for %s" % name
keyfile = os.path.expanduser("~/.ssh/id_dsa")
opts = "StrictHostKeyChecking=no"
sconfig.add_entry(name, host, username, 22, keyfile, opts)

print "Fix known_hosts"
known_hosts = os.path.expanduser('~/.ssh/known_hosts')
host_keys = paramiko.hostkeys.HostKeys(known_hosts)
new_host_keys = paramiko.hostkeys.HostKeys()

ip = socket.gethostbyname(host)

entries = host_keys.lookup(host)
entries = entries._entries if entries else []

ip_entries = host_keys.lookup(ip)
entries += ip_entries._entries if ip_entries else []

for entry in host_keys._entries:
    if entry in entries:
        print "removing %s" % entry
        continue
    new_host_keys._entries.append(entry)

new_host_keys.save(known_hosts)


print "Uploading authorized_keys"

password = getpass.getpass('Enter ssh password: ')

key = open(os.path.expanduser("~/.ssh/id_dsa.pub")).read()
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(host, 22, username, password, allow_agent=False, look_for_keys=False)
print "mkdir -p .ssh"
ssh.exec_command('mkdir -p ~/.ssh/')
print "echo..."
ssh.exec_command('echo "%s" > ~/.ssh/authorized_keys' % key)
print "chmod..."
ssh.exec_command('chmod 644 ~/.ssh/authorized_keys')
ssh.exec_command('chmod 700 ~/.ssh/')
print "Done."