Merge pull request #17 from supertylerc/supertylerc/style-changes

PEP8 Clean Up

SSHClient.py

@@ -4,42 +4,42 @@
 #
 
 # Meta
-__license__ = "AGPLv3"
-__author__ = 'Ahmed Nazmy <[email protected]>'
-
 import getpass
 import logging
-import paramiko
+import os
+import select
+import signal
 import socket
-import tty
 import sys
 import termios
-import signal
-import select
-import os
+import tty
+
+import paramiko
+
+__license__ = "AGPLv3"
+__author__ = 'Ahmed Nazmy <[email protected]>'
 
 TIME_OUT = 10
 
+
 class Client(object):
     def __init__(self, session):
         self._session = session
 
 
-
 class SSHClient(Client):
     def __init__(self, session):
         super(SSHClient, self).__init__(session)
         self._socket = None
         logging.debug("Client: Client Created")
-        
+
     def connect(self, ip, port, size):
         self._size = size
         self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self._socket.settimeout(TIME_OUT)
         self._socket.connect((ip, port))
-        logging.debug("SSHClient: Connected to {0}:{1}".format(ip,port))
-
-
+        logging.debug("SSHClient: Connected to {0}:{1}".format(ip, port))
+
     def get_transport(self):
         transport = paramiko.Transport(self._socket)
         transport.set_keepalive(10)
@@ -61,7 +61,7 @@ def start_session(self, user, auth_secret):
                     transport.auth_password(user, getpass.getpass())
             self._start_session(transport)
         except EOFError as exc:
-            logging.error('Received EOFError.  Assuming bad SSH implementation.')
+            logging.error('EOFError.  Assuming bad SSH implementation.')
             logging.error('Original Erorr: %s', exc)
             self._handle_exception(transport)
         except Exception as e:
@@ -75,7 +75,6 @@ def _handle_exception(self, transport):
             transport.close()
         self._socket.close()
 
-
     def _start_session(self, transport):
         chan = transport.open_session()
         cols, rows = self._size
@@ -87,7 +86,6 @@ def _start_session(self, transport):
         transport.close()
         self._socket.close()
 
-
     def interactive_shell(self, chan):
         # Handle session IO
         sys.stdout.flush()

aker.py

@@ -7,37 +7,36 @@
 
 
 # Meta
+from configparser import ConfigParser
+import getpass
+import logging
+import os
+import time
+import uuid
+
+import paramiko
+
+import tui
+from session import SSHSession
+from snoop import Sniffer
+
 __version__ = '0.2.1'
 __version_info__ = (0, 2, 1)
 __license__ = "AGPLv3"
 __license_info__ = {
     "AGPLv3": {
         "product": "aker",
-        "users": 0, # 0 being unlimited
+        "users": 0,  # 0 being unlimited
         "customer": "Unsupported",
         "version": __version__,
         "license_format": "1.0",
     }
 }
 
-import logging
-import os
-import sys
-import uuid
-import getpass
-import paramiko
-import socket
-from configparser import ConfigParser
-import time
-
-import tui
-from session import SSHSession
-from snoop import Sniffer
-
-
 config_file = "/etc/aker.ini"
 log_file = 'aker.log'
-logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',filename=log_file,level=logging.INFO)
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s',
+                    filename=log_file, level=logging.INFO)
 
 
 class Configuration(object):
@@ -46,91 +45,97 @@ def __init__(self, filename):
         self.src_ip = remote_connection.split()[0]
         self.src_port = remote_connection.split()[1]
         self.session_uuid = uuid.uuid1()
-        #TODO: Check file existance , handle exception
+        # TODO: Check file existance , handle exception
         configparser = ConfigParser()
         if filename:
             configparser.read(filename)
             self.log_level = configparser.get('General', 'log_level')
 
 
-
-
 class User(object):
-    def __init__(self,username):
+    def __init__(self, username):
         self.name = username
         configparser = ConfigParser()
         configparser.read(config_file)
-        #TODO: Add excpetion to handle problems with users in config
+        # TODO: Add excpetion to handle problems with users in config
         hosts = configparser.get(self.name, 'hosts').split("\n")
         self.enabled = configparser.get(self.name, 'enabled')
         self.ssh_hosts = {}
         self.load_ssh_hosts(hosts)
 
-
-    def load_ssh_hosts(self,hosts):
+    def load_ssh_hosts(self, hosts):
         for host in hosts:
-            #TODO: handle exception for incomplete or misplaced entry, i.e host,user,port
-            hostname,port,username = host.split(",")
-            self.ssh_hosts[hostname] = {'port': port , 'username': username}
+            # TODO: handle exception for incomplete or misplaced entry,
+            #       i.e host,user,port
+            hostname, port, username = host.split(",")
+            self.ssh_hosts[hostname] = {'port': port, 'username': username}
 
     def get_priv_key(self):
-        try :
-            #TODO: check better identity options
-            #TODO: get user priv key from configfile
-            privkey = paramiko.RSAKey.from_private_key_file(os.path.expanduser("~/.ssh/id_rsa"))
+        try:
+            # TODO: check better identity options
+            # TODO: get user priv key from configfile
+            privkey = paramiko.RSAKey.from_private_key_file(
+                os.path.expanduser("~/.ssh/id_rsa")
+            )
         except Exception as e:
-            logging.error("Core: Invalid Private Key for user {0} : {1} ".format(self.name, e.message))
+            logging.error("Core: Invalid Private Key for user {0} : {1} "
+                          .format(self.name, e.message))
             raise Exception("Core: Invalid Private Key")
-        else :
+        else:
             return privkey
 
 
 class Aker(object):
     """ Aker core module, this is the management module
     """
 
-    def __init__(self,log_level = 'INFO'):
+    def __init__(self, log_level='INFO'):
         self.config = Configuration(config_file)
         self.posix_user = getpass.getuser()
         self.user = User(self.posix_user)
         self.log_level = self.config.log_level
         self.sniffer = Sniffer()
-        logging.info("Core: Starting up, user={0} from={1}:{2}".format(self.posix_user,self.config.src_ip,self.config.src_port))
+        logging.info("Core: Starting up, user={0} from={1}:{2}".format(
+            self.posix_user, self.config.src_ip, self.config.src_port)
+        )
         self.build_tui()
 
-
-
     def build_tui(self):
         logging.debug("Core: Drawing TUI")
         self.tui = tui.Window(self)
         self.tui.draw()
         self.tui.start()
 
-    def init_connection(self,host):
+    def init_connection(self, host):
         screen_size = self.tui.loop.screen.get_cols_rows()
         logging.debug("Core: pausing TUI")
         self.tui.pause()
         session_uuid = uuid.uuid4()
         session_start_time = time.strftime("%Y%m%d-%H%M%S")
-        session = SSHSession(self,host,session_uuid)
-        #TODO: add err handling
+        session = SSHSession(self, host, session_uuid)
+        # TODO: add err handling
         session.connect(screen_size)
-        #TODO enhance sniffer code
-        session_log_filename = "{0}-{1}-{2}_{3}_{4}.log".format(self.posix_user, host, session_start_time, self.config.src_port, session_uuid)
-        logging.info("Core: Started session UUID {0} for user {1} to host {2}".format(session_uuid,self.posix_user,host))
+        # TODO enhance sniffer code
+        session_log_filename = "{0}-{1}-{2}_{3}_{4}.log".format(
+            self.posix_user, host,
+            session_start_time, self.config.src_port,
+            session_uuid
+        )
+        logging.info("Core: Started session UUID {0} for user {1} to host {2}"
+                     .format(session_uuid, self.posix_user, host))
         self.sniffer.set_log_filename(session_log_filename)
         self.sniffer.capture()
         try:
             session.start_session()
         finally:
             self.sniffer.restore()
             self.tui.restore()
-            #TODO: better handling through session_ended_handler
-            self.tui.search_edit.set_edit_text("") # Clear selected hosts
-
+            # TODO: better handling through session_ended_handler
+            self.tui.search_edit.set_edit_text("")  # Clear selected hosts
 
     def session_end_callback(self, session):
-        logging.info("Core: Finished session UUID {0} for user {1} to host {2}".format(session.uuid,self.posix_user,session.host))
+        logging.info("Core: Finished session UUID {0} for user {1} to host {2}"
+                     .format(session.uuid, self.posix_user, session.host))
 
 
 if __name__ == '__main__':

session.py

@@ -2,59 +2,58 @@
 #
 #       Copyright 2016 Ahmed Nazmy
 #
-
-# Meta
-__license__ = "AGPLv3"
-__author__ = 'Ahmed Nazmy <[email protected]>'
-
 import getpass
 import logging
-import signal
 import os
+import signal
+
 from SSHClient import SSHClient
 
+__license__ = "AGPLv3"
+__author__ = 'Ahmed Nazmy <[email protected]>'
+
+
 class Session(object):
     """ Base class for sessions
         Different type of sessions to be
         added later
     """
-    
-    def __init__(self,aker_core,host,uuid):
-        self.aker= aker_core
+
+    def __init__(self, aker_core, host, uuid):
+        self.aker = aker_core
         self.host = host
         self.host_user = self.aker.user.ssh_hosts[host]['username']
         self.host_port = int(self.aker.user.ssh_hosts[host]['port'])
         self.uuid = uuid
         logging.debug("Session: Base Session created")
-        #TODO : Add UUID shit 
-        
+        # TODO : Add UUID shit
+
     def connect(self, size):
         self._client.connect(self.host, self.host_port, size)
-        
+
     def start_session(self):
         raise NotImplementedError
-        
+
     def close_session(self):
         self.aker.session_end_callback(self)
-        
+
     def kill_session(self, signum, stack):
-        #TODO : Change behavoir to show screen again
+        # TODO : Change behavoir to show screen again
         logging.debug("Session: Session Killed")
         self.close_session()
         os.kill(os.getpid(), signal.SIGKILL)
 
 
 class SSHSession(Session):
-    """ Wrapper around SSHClient instantiating 
+    """ Wrapper around SSHClient instantiating
         a new SSHClient instance everytime
     """
-    
-    def __init__(self, aker_core, host,uuid):
-        super(SSHSession, self).__init__(aker_core, host,uuid)
+
+    def __init__(self, aker_core, host, uuid):
+        super(SSHSession, self).__init__(aker_core, host, uuid)
         self._client = SSHClient(self)
         logging.debug("Session: SSHSession created")
 
-
     def start_session(self):
         try:
             auth_secret = self.aker.user.get_priv_key()

snoop.py

@@ -1,21 +1,16 @@
 # -*- coding: utf-8 -*-
 #
-#       Copyright 2016 Ahmed Nazmy 
+#       Copyright 2016 Ahmed Nazmy
 #
-
-# Meta
-__license__ = "AGPLv3"
-__author__ = 'Ahmed Nazmy <[email protected]>'
-
-
-
-from Queue import Queue
 import logging
+import os
+from Queue import Queue
 import stat
 import threading
-import os
 import sys
-import codecs
+
+__license__ = "AGPLv3"
+__author__ = 'Ahmed Nazmy <[email protected]>'
 
 
 class Sniffer(object):
@@ -57,6 +52,7 @@ def restore(self):
             self._sniffer.join()
             self.log_filename = None
 
+
 class SessionSniffer(threading.Thread):
     def __init__(self, keys_queue):
         threading.Thread.__init__(self)
@@ -71,7 +67,12 @@ def set_log_filename(self, log_filename):
     def run(self):
         if self._log_filename:
             self._log_file = open(self._log_filename, "wb")
-            os.chmod(self._log_file.name, stat.S_IREAD | stat.S_IWRITE | stat.S_IWRITE | stat.S_IRGRP | stat.S_IROTH)
+            os.chmod(self._log_file.name,
+                     stat.S_IREAD |
+                     stat.S_IWRITE |
+                     stat.S_IWRITE |
+                     stat.S_IRGRP |
+                     stat.S_IROTH)
         while True:
             if not self._session_stop:
                 c = self._key_queue.get()