Code cleanup #1

Merged
zlax merged 1 commits from Difrex/master into master 2017-07-25 13:48:14 +00:00
1 changed files with 174 additions and 144 deletions

View File

@ -4,20 +4,21 @@
from __future__ import division
from os import path
from sys import argv
from math import *
from math import ceil
from getpass import getpass
from optparse import OptionParser
from curses import textpad, panel
from curses import textpad
import ConfigParser
import subprocess
import base64
import curses
# https://github.com/zlaxy/sshch
version="0.55"
version = "0.55"
# 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)
@ -26,21 +27,25 @@ def AddNewAlias(alias):
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 = ""
@ -51,47 +56,52 @@ def ConnectAlias(alias, command=False):
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]
# Variables bellow is newer used
subprocess.Popen(exec_string, shell=True)
# 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")
if result:
prompt_add = ("".join(["Enter connection string for new alias ",
"(example: ssh user@somehost.com):\n"]))
string = ""
while string == "":
string = raw_input (prompt_add)
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")
prompt_edit = ("".join(["Enter connection string for existing alias ",
"(example: ssh user@somehost.com):\n"]))
string = ""
while string == "":
string = raw_input (prompt_edit)
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)
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)
prompt_remove = ("Type 'yes' if you sure to remove '" + alias + "' alias: ")
string = raw_input(prompt_remove)
if string == "yes":
RemoveAliases([alias])
else:
@ -99,6 +109,7 @@ def CMDRemove(alias):
else:
print "error: '" + alias + "' alias is not exists."
def CMDConnect(aliases, command=False):
for alias in aliases:
if conf.has_section(alias):
@ -106,13 +117,17 @@ def CMDConnect(aliases, command=False):
else:
print "error: '" + alias + "' alias is not exists"
def CMDList(option, opt, value, parser):
print ', '.join(str(p) for p in conf.sections())
def CMDFullList(option, opt, value, parser):
print '\n'.join((str(p) + " - " + conf.get(p, "exec_string") +
(" [password]" if conf.has_option(p, "password") else ""))
for p in conf.sections())
for p in conf.sections():
to_print = "".join([str(p), " - ", conf.get(p, "exec_string"),
(" [password]" if conf.has_option(p, "password") else ""), "\n"])
print(to_print)
def CursesConnect(screen, aliases, command=False):
curses.endwin()
@ -121,21 +136,25 @@ def CursesConnect(screen, aliases, command=False):
print "Press 'enter' to continue."
screen.getch()
def CursesExit(error=False):
curses.endwin()
if error: print error
if error:
print error
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.addnstr(0, 0, title, w , text_colorpair)
title_window = new_window.subwin(1, w, y, x)
title_window.addnstr(0, 0, title, w, text_colorpair)
title_window.refresh()
sub_window = new_window.subwin(h, w, y + 1, x)
textbox_field = textpad.Textbox(sub_window, insert_mode=True)
@ -143,10 +162,11 @@ def CursesTextpad(screen, h, w, y, x, title="", value="",
new_window.box()
new_window.attroff(deco_colorpair)
new_window.refresh()
sub_window.addnstr(0, 0 , value, h * w, text_colorpair)
sub_window.addnstr(0, 0, value, h * w, 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)
@ -154,9 +174,9 @@ def CursesPanel(screen, h, w, y, x, text,
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 = new_window.subwin(h - 2, w - 2, y + 1, x + 1)
sub_window.insstr(0, 0, text)
panel = curses.panel.new_panel(new_window)
curses.panel.new_panel(new_window)
curses.panel.update_panels()
screen.refresh()
if confirm == "password":
@ -193,28 +213,30 @@ def CursesPanel(screen, h, w, y, x, text,
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")
epilog = ("".join(["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",
opts.add_option('-l', '--list', action="callback",
callback=CMDList, help="show list of all existing aliases")
opts.add_option('-f', '--fulllist', action = "callback",
opts.add_option('-f', '--fulllist', action="callback",
callback=CMDFullList, help=("show list of all existing " +
"aliases with connection strings"))
opts.add_option('-a', '--add', action="store", type="string",
@ -233,27 +255,34 @@ def CMDOptions():
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)
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)
help_screen = ("".join([" 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)]
@ -278,21 +307,20 @@ def CursesMain():
page = 1
for i in range(1, max_row + 1):
if row_num == 0:
box.addnstr(1, 1, "There aren't any aliases yet. Press 'a'" +
" to add new one.", width - 6, highlight_text)
box.addnstr(1, 1, "There aren't any aliases yet. Press 'a' to add new one.",
width - 6, highlight_text)
else:
if conf.has_option(strings[i - 1], "password"):
password = " [password]"
else:
password = ""
exec_string = ["[", selected_strings[i], "] ", str(i), " ",
strings[i - 1], " (", conf.get(strings[i - 1], "exec_string"),
")", password]
if (i == position):
box.addnstr(i, 2, "[" + selected_strings[i] + "] " +
str(i) + " " + strings[i - 1] + " (" +
conf.get(strings[i - 1], "exec_string") + ")" +
(" [password]" if conf.has_option(strings[i - 1],
"password") else ""), width - 6, highlight_text)
box.addnstr(i, 2, "".join(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") + ")" +
(" [password]" if conf.has_option(strings[i - 1],
"password") else ""), width - 6, normal_text)
box.addnstr(i, 2, "".join(exec_string), width - 6, normal_text)
if i == row_num:
break
screen.refresh()
@ -315,7 +343,8 @@ def CursesMain():
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,
if not add_result:
CursesPanel(screen, 3,
width - 6, (height // 2) - 1, 3, add_result,
normal_text, highlight_text)
else:
@ -351,8 +380,8 @@ def CursesMain():
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")
" 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(
@ -363,11 +392,12 @@ def CursesMain():
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)")
remove_confirm = (
"".join(["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)")
remove_confirm = ("".join(["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,
@ -391,8 +421,10 @@ def CursesMain():
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)
"".join([
"Enter specific command to execute with selected ",
"alias/aliases:"]
), "", normal_text, highlight_text)
command_string = command_textpad.edit(CursesTextpadConfirm)
CursesConnect(screen, selected,
command_string.replace("\n", "").rstrip())
@ -409,7 +441,8 @@ def CursesMain():
curses.KEY_IC)) and row_num != 0:
if selected_strings[position] == ' ':
selected_strings[position] = '*'
else: selected_strings[position] = ' '
else:
selected_strings[position] = ' '
if page == 1:
if position < i:
position = position + 1
@ -471,25 +504,22 @@ def CursesMain():
for i in range(1 + (max_row * (page - 1)), max_row + 1 +
(max_row * (page - 1))):
if row_num == 0:
box.addnstr(1, 1, "There aren't any aliases yet. " +
"Press 'a' to add new one.", width - 6,
highlight_text)
box.addnstr(1, 1, "There aren't any aliases yet. Press 'a' to add new one.",
width - 6, 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") + ")" + (" [password]" if
conf.has_option(strings[i - 1], "password") else
""), width - 6, highlight_text)
if conf.has_option(strings[i - 1], "password"):
password = " [password]"
else:
box.addnstr(i - (max_row * (page - 1)), 2, "[" +
selected_strings[i] + "] " + str(i) + " " +
strings[i - 1] + " (" + conf.get(strings[i - 1],
"exec_string") + ")" + (" [password]" if
conf.has_option(strings[i - 1], "password") else
""), width - 6, normal_text)
password = ""
exec_string = ["[", selected_strings[i], "] ", str(i), " ",
strings[i - 1], " (", conf.get(strings[i - 1], "exec_string"),
")", password]
if (i + (max_row * (page - 1)) == (position + (max_row * (page - 1)))):
box.addnstr(i - (max_row * (page - 1)), 2, "".join(
exec_string), width - 6, highlight_text)
else:
box.addnstr(i - (max_row * (page - 1)), 2, "".join(
exec_string), width - 6, normal_text)
if i == row_num:
break
screen.refresh()
@ -497,6 +527,7 @@ def CursesMain():
key_pressed = screen.getch()
CursesExit()
if __name__ == "__main__":
conf = ConfigParser.RawConfigParser()
if not path.exists(conf_file):
@ -508,8 +539,7 @@ if __name__ == "__main__":
except KeyboardInterrupt:
exit()
except ConfigParser.Error:
print ("Error: can't parse your config file, please check" +
" it manually or make new one")
print ("Error: can't parse your config file, please check it manually or make new one")
exit()
else:
try:
@ -517,8 +547,8 @@ if __name__ == "__main__":
except KeyboardInterrupt:
CursesExit()
except ConfigParser.NoOptionError:
CursesExit("Error: can't parse your config file, please " +
"check it manually or make new one")
CursesExit("".join(["Error: can't parse your config file, please ",
"check it manually or make new one"]))
except curses.error:
CursesExit("Error: can't show some curses element, maybe " +
"your terminal is too small")
CursesExit("".join(["Error: can't show some curses element, maybe ",
"your terminal is too small"]))