#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: Wu Liang
@contact: garcia.wul@alibaba-inc.com
@date: 13-8-8
@version: 0.0.0
"""

from argparse import ArgumentParser
import getpass
import os
import socket
import sys
import urllib2

try:
    import termios
    import tty
    has_termios = True
except ImportError:
    has_termios = False

import demjson
from mako.template import Template
import paramiko
import re
import yaml
import tempfile

# 参考：http://stackoverflow.com/questions/566746/how-to-get-console-window-width-in-python
def getTerminalSize():
    import os
    env = os.environ
    def ioctl_GWINSZ(fd):
        try:
            import fcntl, termios, struct, os
            cr = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ,
        '1234'))
        except:
            return
        return cr
    cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2)
    if not cr:
        try:
            fd = os.open(os.ctermid(), os.O_RDONLY)
            cr = ioctl_GWINSZ(fd)
            os.close(fd)
        except:
            pass
    if not cr:
        cr = (env.get('LINES', 25), env.get('COLUMNS', 80))
    return int(cr[1]), int(cr[0])

# copy from paramiko
def interactive_shell(chan):
    if has_termios:
        posix_shell(chan)
    else:
        windows_shell(chan)

def posix_shell(chan):
    import select
    oldtty = termios.tcgetattr(sys.stdin)
    try:
        tty.setraw(sys.stdin.fileno())
        tty.setcbreak(sys.stdin.fileno())
        chan.settimeout(0.0)

        while True:
            r, w, e = select.select([chan, sys.stdin], [], [])
            if chan in r:
                try:
                    x = chan.recv(1024)
                    if len(x) == 0:
                        print '\r\n*** EOF\r\n',
                        break
                    sys.stdout.write(x)
                    sys.stdout.flush()
                except socket.timeout:
                    pass
            if sys.stdin in r:
                x = sys.stdin.read(1)
                if len(x) == 0:
                    break
                chan.send(x)

    finally:
        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)

# thanks to Mike Looijmans for this code
def windows_shell(chan):
    import threading

    sys.stdout.write("Line-buffered terminal emulation. Press F6 or ^Z to send EOF.\r\n\r\n")

    def writeall(sock):
        while True:
            data = sock.recv(256)
            if not data:
                sys.stdout.write('\r\n*** EOF ***\r\n\r\n')
                sys.stdout.flush()
                break
            sys.stdout.write(data)
            sys.stdout.flush()

    writer = threading.Thread(target=writeall, args=(chan,))
    writer.start()

    try:
        while True:
            d = sys.stdin.read(1)
            if not d:
                break
            chan.send(d)
    except EOFError:
        # user hit ^Z or F6
        pass

def login(host, port, username, password):
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(host, port, username, password)
    width, height = getTerminalSize()
    channel = client.invoke_shell(term="xterm", width=width, height=height)
    interactive_shell(channel)
    channel.close()
    client.close()

def loginByExpect(host, port, username, password):
    loginScript = """#!/usr/bin/expect -f
spawn ssh -l ${username} ${host} -X -Y
expect {
    "*Are you sure you want to continue connecting*" {
    send "yes\\r"
    expect "*assword*"
    send "${password}\\r"
}
    "*assword*" {
    send "${password}\\r"
    }
    "*$*" {}
}
interact
"""
    script = Template(loginScript).render(
        username = username,
        host = host,
        password =  password
    )
    loginScriptFileInstance = tempfile.NamedTemporaryFile(delete=False)
    loginScriptFilename = loginScriptFileInstance.name
    loginScriptFileStream = file(loginScriptFilename, "w")
    loginScriptFileStream.write(script)
    loginScriptFileStream.close()
    os.chmod(loginScriptFilename, 0777)
    loginScriptFileInstance.close()
    os.system(loginScriptFilename)
    try:
        os.remove(loginScriptFilename)
    except Exception:
        pass

def showItems(items, config):
    newItems = []
    index = 1
    for item in items:
        item["index"] = index
        newItems.append(item)
        index += 1
    itemsTemplate = Template('''% for item in items:
NO.${item["index"]}\t${item["username"]}@${item["host"]}
% endfor
    ''')
    print itemsTemplate.render(items=newItems)
    choice = raw_input("Your choice:\t")
    choice = choice.strip()
    if not choice.isdigit():
        return
    choice = int(choice)
    if choice > len(items):
        return
    hostData = items[choice - 1]
    if config.get("LoginMode", "paramiko") == "expect":
        loginByExpect(hostData["host"],
                      hostData["port"],
                      hostData["username"],
                      hostData["password"]
        )
    else:
        login(hostData["host"],
              hostData["port"],
              hostData["username"],
              hostData["password"]
        )

def getUserDataFromServiceApi(config):
    url = Template("${serviceApi}/get/").render(
        serviceApi = config["ServiceApi"].rstrip("/")
    )
    content = urllib2.urlopen(url).read()
    content = demjson.decode(content, strict=False, encoding="utf-8")
    showItems(content, config)

def getUserDataFromConfigFile(config):
    content = config["hosts"]
    showItems(content, config)

def loginManually(host, config):
    pattern1 = re.compile("(.*)@(.*)")
    matched = pattern1.search(host)
    if matched:
        username, host = matched.groups()
    else:
        username = getpass.getuser()
        host = host
    password = getpass.getpass(Template("${username}@${host}'s password: ").render(
        username = username,
        host = host
    ))
    if config.get("AutoSave", False):
        url = Template("${serviceApi}/add/${host}/${port}/${username}/${password}/").render(
            serviceApi = config["ServiceApi"],
            host = host,
            port = 22,
            username = username,
            password = password
        )
        urllib2.urlopen(url)
    login(host, 22, username, password)

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument("--config", dest="configFile",
        default=os.path.join(os.getenv("HOME"), ".xlogin", "xlogin.yaml"))
    parser.add_argument("-l", "--login", dest="login")
    argument = parser.parse_args()

    if not os.path.exists(argument.configFile):
        print >> sys.stderr, "%s does not exist" % argument.configFile
        sys.exit(1)
    config = yaml.load(open(argument.configFile))
    if argument.login:
        loginManually(argument.login, config)
    elif config.has_key("ServiceApi"):
        getUserDataFromServiceApi(config)
    else:
        getUserDataFromConfigFile(config)
