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

from __future__ import division
from os import path
from sys import argv
from math import *
from getpass import getpass
from optparse import OptionParser
from curses import textpad, panel
import ConfigParser
import subprocess
import base64
import curses

# https://github.com/zlaxy/sshch
version="0.4"
# path to conf file, default: ~/.config/sshch.conf
conf_file = path.expanduser("~") + '/.config/sshch.conf'

def AddNewAlias(alias):
    if not conf.has_section(alias):
        conf.add_section(alias)
        conf.write(open(conf_file, "w"))
        return True
    else:
        return "error: '" + alias + "' already exists"

def SetAliasString(alias, string):
    conf.set(alias, "exec_string", string)
    conf.write(open(conf_file, "w"))

def SetPassword(alias, string):
    string = base64.b64encode(base64.b16encode(
                              base64.b32encode(string)))
    conf.set(alias, "password", string)
    conf.write(open(conf_file, "w"))

def RemoveAliases(aliases):
    for alias in aliases:
        conf.remove_section(alias)
        conf.write(open(conf_file, "w"))

def ConnectAlias(alias, command=False):
    print "Connecting to " + alias + "..."
    exec_string = ""
    if conf.has_option(alias, "password"):
        password = base64.b32decode(base64.b16decode(
            base64.b64decode(conf.get(alias, "password"))))
        exec_string = "sshpass -p " + password + " "
    exec_string = exec_string + conf.get(alias, "exec_string")
    if command:
        exec_string = exec_string + " " + command
    p = subprocess.Popen(exec_string, shell=True)
    streamdata = p.communicate()[0]
    print "... " + alias + " session finished."

def CMDAdd(alias):
    result = AddNewAlias(alias)
    if result == True:
        prompt_add = ("Enter connection string for new alias " +
                      "(example: ssh user@somehost.com):\n")
        string = ""
        while string == "":
            string = raw_input (prompt_add)
        SetAliasString(alias, string)
    else:
        print result

def CMDEdit(alias):
    if conf.has_section(alias):
        prompt_edit = ("Enter connection string for existing alias " +
                       "(example: ssh user@somehost.com):\n")
        string = ""
        while string == "":
            string = raw_input (prompt_edit)
        SetAliasString(alias, string)
    else:
        print "error: '" + alias + "' alias is not exists"

def CMDPassword(alias):
    if conf.has_section(alias):
        prompt_pass = ("[UNSAFE] Enter password for sshpass: ")
        string = ""
        string = getpass (prompt_pass)
        if not string == "": SetPassword(alias, string)
    else:
        print "error: '" + alias + "' alias is not exists"

def CMDRemove(alias):
    if conf.has_section(alias):
        prompt_remove = ("Type 'yes' if you sure to remove '" +
                         alias + "' alias: ")
        string = raw_input (prompt_remove)
        if string == "yes":
            RemoveAliases([alias])
        else:
            print "'" + alias + "' alias was not deleted."
    else:
        print "error: '" + alias + "' alias is not exists."

def CMDConnect(aliases, command=False):
    for alias in aliases:
        if conf.has_section(alias):
            ConnectAlias(alias, command)
        else:
            print "error: '" + alias + "' alias is not exists"

def CMDList(option, opt, value, parser):
    print ', '.join(str(p) for p in conf.sections())

def CursesConnect(screen, aliases, command=False):
    curses.endwin()
    for alias in aliases:
        ConnectAlias(alias, command)
    print "Press 'enter' to continue."
    screen.getch()

def CursesExit():
    curses.endwin()
    exit()

def CursesTextpadConfirm(value):
    if value == 10:
        value = 7
    return value

def CursesTextpad(screen, h, w, y, x, title="", value="",
                  text_colorpair=0, deco_colorpair=0):
    new_window = curses.newwin(h + 3, w + 2, y - 1, x - 1)
    title_window = new_window.subwin(1, w , y , x)
    title_window.addstr(0, 0, title, text_colorpair)
    title_window.refresh()
    sub_window = new_window.subwin(h, w, y + 1, x)
    textbox_field = textpad.Textbox(sub_window, insert_mode=True)
    new_window.attron(deco_colorpair)
    new_window.box()
    new_window.attroff(deco_colorpair)
    new_window.refresh()
    sub_window.addstr(0, 0 , value, text_colorpair)
    sub_window.attron(text_colorpair)
    return textbox_field

def CursesPanel(screen, h, w, y, x, text,
                text_colorpair=0, deco_colorpair=0, confirm=0):
    new_window = curses.newwin(h, w, y, x)
    new_window.erase()
    new_window.attron(deco_colorpair)
    new_window.box()
    new_window.attroff(deco_colorpair)
    sub_window = new_window.subwin(h - 2, w - 2 , y + 1 , x + 1 )
    sub_window.addstr(0, 0, text)
    panel = curses.panel.new_panel(new_window)
    curses.panel.update_panels()
    screen.refresh()
    if confirm == "password":
        hidden_password = ""
        keych = ""
        position = 2
        while True:
            keych = screen.getch()
            if keych == ord("\n"):
                break
            if keych == 27:
                hidden_password = ""
                break
            if keych == curses.KEY_BACKSPACE:
                if position > 2:
                    sub_window.addstr(1, position - 1, " ",
                                      text_colorpair)
                    sub_window.refresh()
                    position = position - 1
                    hidden_password = hidden_password[0:-1]
            else:
                hidden_password += curses.keyname(keych)
                sub_window.addstr(1, position, "*", text_colorpair)
                position += 1
                sub_window.refresh()
        return hidden_password
    elif confirm == "remove":
        keych = screen.getch()
        if keych == ord("y") or keych == ord("Y"):
            return "confirm"
        return False
    else:
        screen.getch()

def CMDOptions():
    class FormatedParser(OptionParser):
        def format_epilog(self, formatter):
            return self.epilog
    usage = "usage: %prog [options] [aliases]"
    progname = path.basename(__file__)
    epilog = ("Examples:\n  " + 
        progname + " existingalias\n  " +
        progname + " -a newremoteserver\n  " + 
        progname + " --edit=newremoteserver -p newremoteserver\n  " +
        progname + ' -c "ls -l" newremoteserver\n  ' + 
        progname + " -c reboot existingalias newremoteserver\n"
        "Examples of connection string:\n  " +
        "ssh user@somehost.com\n  " +
        "ssh gates@8.8.8.8 -p 667\n  " +
        "ssh root@somehost.com -t tmux a\n" +
        "Also, you can edit config file manually: " + conf_file + "\n")
    opts = FormatedParser(usage=usage, version="%prog " + version,
                          epilog=epilog)
    opts.add_option('-l', '--list', action = "callback",
        callback=CMDList, help="show list of all existing aliases")
    opts.add_option('-a', '--add', action="store", type="string",
        dest="add", metavar="alias", default=False,
        help="add new alias for connection string")
    opts.add_option('-c', '--command', action="store", type="string",
        dest="command", metavar="command", default=False,
        help="add command for executing alias")
    opts.add_option('-e', '--edit', action="store", type="string",
        dest='edit', metavar="alias", default=False,
        help="edit existing connection string")
    opts.add_option('-p', '--password', action="store", type="string",
        dest='password', metavar="alias", default=False,
        help="set and store password for sshpass [UNSAFE]")
    opts.add_option('-r', '--remove', action="store", type="string",
        dest='remove', metavar="alias", default=False,
        help="remove existing alias of connection string")
    options, alias = opts.parse_args()
    if options.add: CMDAdd(options.add)
    if options.edit: CMDEdit(options.edit)
    if options.password: CMDPassword(options.password)
    if options.remove: CMDRemove(options.remove)
    if alias: CMDConnect(alias, options.command)

# curses template from: https://stackoverflow.com/a/30828805/6224462
def CursesMain():
    help_screen = (" Press:\n" +
    "  'z'/'x' or arrows - navigation\n" +
    "  'a'/'F2' - add new alias\n" +
    "  'e'/'F4' - edit existing alias\n" +
    "  'p'/'F6' - set alias's password for sshpass [UNSAFE]\n" +
    "  'space'/'insert' - select\n" +
    "  'r'/'F8' - remove selected alias/aliases\n" +
    "  'c'/'F3' - execute specific command with selected alias/aliases\n" +
    "  'enter'/'F9' - connect to selected alias/aliases\n" +
    "  'q'/'F10' - quit\n" +
    " Run program with '--help' option to view command line help.\n" +
    " Also, you can edit config file manually:\n" +
    "  " + conf_file)
    strings = conf.sections()
    row_num = len(strings)
    selected_strings = [" " for i in range(0, row_num + 1)]
    screen = curses.initscr()
    curses.noecho()
    curses.cbreak()
    curses.start_color()
    screen.keypad(1)
    curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_CYAN)
    highlight_text = curses.color_pair(1)
    normal_text = curses.A_NORMAL
    height, width = screen.getmaxyx()
    screen.border(0)
    curses.curs_set(0)
    max_row = height - 5
    screen.addstr(1, 2, "sshch " + version + ", press 'h' for help")
    box = curses.newwin(max_row + 2, width - 2, 2, 1)
    box.box()
    pages = int(ceil(row_num / max_row))
    position = 1
    page = 1
    for i in range(1, max_row + 1):
        if row_num == 0:
            box.addstr(1, 1, "There aren't any aliases yet. Press 'a'" +
                       " to add new one.", highlight_text)
        else:
            if (i == position):
                box.addnstr(i, 2, "[" + selected_strings[i] + "] " +
                    str(i) + " " + strings[i - 1] + " (" +
                    conf.get(strings[i - 1], "exec_string") + ")",
                    width - 6, highlight_text)
            else:
                box.addnstr(i, 2, "[" + selected_strings[i] + "] " +
                    str(i) + " " + strings[i - 1] + " (" +
                    conf.get(strings[i - 1], "exec_string") + ")",
                    width - 6, normal_text)
            if i == row_num:
                break
    screen.refresh()
    box.refresh()
    key_pressed = screen.getch()
    while True:
        if key_pressed == ord('q') or key_pressed == ord(
                              'Q') or key_pressed == (
                               27) or key_pressed == curses.KEY_F10:
            CursesExit()
        if key_pressed == ord('h') or key_pressed == ord(
                              'H') or key_pressed == curses.KEY_F1:
            CursesPanel(screen, height - 4, width - 6, 2, 3,
                        help_screen, normal_text, highlight_text)
        if key_pressed == ord('a') or key_pressed == ord(
                              'A') or key_pressed == curses.KEY_F2:
            new_alias_textpad = CursesTextpad(screen, 1, width - 8,
                (height // 2) - 1, 4, "Enter new alias:", "",
                normal_text, highlight_text)
            add_alias = new_alias_textpad.edit(CursesTextpadConfirm)
            if not add_alias.rstrip() == "":
                add_result = AddNewAlias(add_alias.rstrip())
                if not add_result == True: CursesPanel(screen, 3,
                    width - 6, (height // 2) - 1, 3, add_result,
                    normal_text, highlight_text)
                else:
                    add_string = ""
                    while add_string.rstrip() == "":
                        string_textpad = CursesTextpad(screen, 3,
                            width - 8, (height // 2) - 1, 4,
                            "Enter full execution string:",
                            "ssh ", normal_text, highlight_text)
                        add_string = string_textpad.edit(
                            CursesTextpadConfirm)
                    SetAliasString(add_alias.rstrip(),
                                   add_string.rstrip())
                    strings = conf.sections()
                    row_num = len(strings)
                    selected_strings.append(" ")
                    pages = int(ceil(row_num / max_row))
                    box.refresh()
        if (key_pressed == ord('e') or key_pressed == ord(
                'E') or key_pressed == curses.KEY_F4) and row_num != 0:
            edit_string = ""
            while edit_string.rstrip() == "":
                string_textpad = CursesTextpad(screen, 3, width - 8,
                    (height // 2) - 1, 4, "Enter new execution string:",
                    conf.get(strings[position - 1], "exec_string"),
                    normal_text, highlight_text)
                edit_string = string_textpad.edit(CursesTextpadConfirm)
            SetAliasString(strings[position - 1], edit_string.rstrip())
            strings = conf.sections()
        if (key_pressed == ord('p') or key_pressed == ord(
                'P') or key_pressed == curses.KEY_F6) and row_num != 0:
            set_password = ""
            set_password = CursesPanel(screen, 4, width - 6,
                (height // 2) - 1, 3,
                " Enter user password for sshpass and press 'enter':" +
                "\n>", normal_text, highlight_text, "password")
            if not set_password == "":
                SetPassword(strings[position - 1], set_password)
        if (key_pressed == ord('r') or key_pressed == ord(
                'R') or key_pressed == curses.KEY_F8 or key_pressed == (
                curses.KEY_DC)) and row_num != 0:
            selected = []
            for i in range(1, row_num + 1):
                if selected_strings[i] == "*":
                    selected.append(strings[i - 1])
            if len(selected) > 0:
                remove_confirm = ("Are you sure to remove " +
                    str(len(selected)) + " selected aliases? (y/N)")
            else:
                remove_confirm = ("Are you sure to remove '" +
                    strings[position - 1] + "' alias? (y/N)")
                selected.append(strings[position - 1])
            remove_result = CursesPanel(screen, 4, width - 6,
                (height // 2) - 1, 3, remove_confirm, normal_text,
                highlight_text, "remove")
            if remove_result == "confirm":
                RemoveAliases(selected)	
                strings = conf.sections()
                row_num = len(strings)
                selected_strings = [" " for i in range(0, row_num + 1)]
                pages = int(ceil(row_num / max_row))
                position = 1
                page = 1
                box.refresh()
        if (key_pressed == ord('c') or key_pressed == ord(
                'C') or key_pressed == curses.KEY_F3) and row_num != 0:
            selected = []
            for i in range(1, row_num + 1):
                if selected_strings[i] == "*":
                    selected.append(strings[i - 1])
            if not len(selected) > 0:
                selected.append(strings[position - 1])
            command_textpad = CursesTextpad(screen, 3, width - 8,
                (height // 2) - 1, 4,
                "Enter specific command to execute with selected " +
                "alias/aliases:", "", normal_text, highlight_text)
            command_string = command_textpad.edit(CursesTextpadConfirm)
            CursesConnect(screen, selected, command_string.rstrip())
        if (key_pressed == ord("\n") or key_pressed == (
                curses.KEY_F9)) and row_num != 0:
            selected = []
            for i in range(1, row_num + 1):
                if selected_strings[i] == "*":
                    selected.append(strings[i - 1])
            if not len(selected) > 0:
                selected.append(strings[position - 1])
            CursesConnect(screen, selected)
        if key_pressed == 32 or key_pressed == curses.KEY_IC:
            if selected_strings[position] == ' ':
                selected_strings[position] = '*'
            else: selected_strings[position] = ' '
            if page == 1:
                if position < i:
                    position = position + 1
                else:
                    if pages > 1:
                        page = page + 1
                        position = 1 + (max_row * (page - 1))
            elif page == pages:
                if position < row_num:
                    position = position + 1
            else:
                if position < max_row + (max_row * (page - 1)):
                    position = position + 1
                else:
                    page = page + 1
                    position = 1 + (max_row * (page - 1))
        if key_pressed == curses.KEY_DOWN or key_pressed == ord(
                                     'x') or key_pressed == ord('X'):
            if page == 1:
                if position < i:
                    position = position + 1
                else:
                    if pages > 1:
                        page = page + 1
                        position = 1 + (max_row * (page - 1))
            elif page == pages:
                if position < row_num:
                    position = position + 1
            else:
                if position < max_row + (max_row * (page - 1)):
                    position = position + 1
                else:
                    page = page + 1
                    position = 1 + (max_row * (page - 1))
        if key_pressed == curses.KEY_UP or key_pressed == ord(
                                   'z') or key_pressed == ord('Z'):
            if page == 1:
                if position > 1:
                    position = position - 1
            else:
                if position > (1 + (max_row * (page - 1))):
                    position = position - 1
                else:
                    page = page - 1
                    position = max_row + (max_row * (page - 1))
        if key_pressed == curses.KEY_LEFT or (key_pressed == 
                          curses.KEY_PPAGE):
            if page > 1:
                page = page - 1
                position = 1 + (max_row * (page - 1))
        if key_pressed == curses.KEY_RIGHT or (key_pressed ==
                          curses.KEY_NPAGE):
            if page < pages:
                page = page + 1
                position = (1 + (max_row * (page - 1)))
        box.erase()
        screen.border(0)
        box.border(0)
        for i in range(1 + (max_row * (page - 1)), max_row + 1 +
                       (max_row * (page - 1))):
            if row_num == 0:
                box.addstr(1, 1, "There aren't any aliases yet. Press" +
                           " 'a' to add new one.", highlight_text)
            else:
                if (i + (max_row * (page - 1)) == (position +
                    (max_row * (page - 1)))):
                    box.addnstr(i - (max_row * (page - 1)), 2, "[" +
                        selected_strings[i] + "] " + str(i) + " " +
                        strings[i - 1] + " (" + conf.get(strings[i - 1],
                        "exec_string") + ")", width - 6, highlight_text)
                else:
                    box.addnstr(i - (max_row * (page - 1)), 2, "[" +
                        selected_strings[i] + "] " + str(i) + " " +
                        strings[i - 1] + " (" + conf.get(strings[i - 1],
                        "exec_string") + ")", width - 6, normal_text)
                if i == row_num:
                    break
        screen.refresh()
        box.refresh()
        key_pressed = screen.getch()    
    CursesExit()

if __name__ == "__main__":
    conf = ConfigParser.RawConfigParser()
    if not path.exists(conf_file):
        open(conf_file, 'w')
    conf.read(conf_file)
    if len(argv) > 1:
        CMDOptions()
    else:
        CursesMain()