Added command line support

Author: humiaozuzu

ChangeLog.txt

@@ -1,5 +1,14 @@
 # YaH3C changelog
 
+## 0.5
+
+* add command line argument support
+* dhcp after authentication
+
+## 0.4
+
+* remove plugin module
+
 ## 0.3
 
 * add eap-md5 support

yah3c/eapauth.py

@@ -9,6 +9,7 @@
 
 import socket
 import os, sys, pwd
+from subprocess import call
 
 from colorama import Fore, Style, init
 # init() # required in Windows
@@ -105,8 +106,14 @@ def EAP_handler(self, eap_packet):
         code, id, eap_len = unpack("!BBH", eap_packet[4:8])
         if code == EAP_SUCCESS:
             display_prompt(Fore.YELLOW, 'Got EAP Success')
-            # invoke plugins 
-            daemonize('/dev/null','/tmp/daemon.log','/tmp/daemon.log')
+            
+            if self.login_info['dhcp_command']:
+                display_prompt(Fore.YELLOW, 'Obtaining IP Address:')
+                call([self.login_info['dhcp_command'], self.login_info['ethernet_interface']])
+
+            if self.login_info['daemon'] == 'True':
+                daemonize('/dev/null','/tmp/daemon.log','/tmp/daemon.log')
+        
         elif code == EAP_FAILURE:
             if (self.has_sent_logoff):
                 display_prompt(Fore.YELLOW, 'Logoff Successfully!')

yah3c/usermgr.py

@@ -14,8 +14,11 @@ class UserMgr:
         "username": "maple",
         "password": "valley",
         "ethernet_interface": "eth0",
-        "carry_version_info": true,
-        "broadcast_logoff": false
+        "dhcp_command": "dhcpcd",
+        "daemon": True,
+        # following has not implemented yet
+        "carry_version_info": True,
+        "broadcast_logoff": False
         "packet_type": "unicast"
     }
     """
@@ -26,6 +29,12 @@ def __init__(self, path=None):
             self.users_cfg_path = path
         self.config = ConfigParser.ConfigParser()
         self.config.read(self.users_cfg_path)
+
+    def save_and_reload(self):
+        fp = open(self.users_cfg_path, 'w')
+        self.config.write(fp)
+        fp.close()
+        self.config.read(self.users_cfg_path)
 
     def get_user_number(self):
         return len(self.config.sections())
@@ -40,23 +49,25 @@ def get_all_users_info(self):
         return users_info
 
     def get_user_info(self, username):
-        return dict(self.config.items(username))
+        user_info = dict(self.config.items(username))
+        user_info['username'] = username
+        return user_info
 
     def add_user(self, user_info):
         self.config.add_section(user_info['username'])
         self.update_user_info(user_info)
 
     def remove_user(self, username):
         self.config.remove_section(username)
-        fp = open(self.users_cfg_path, 'w')
-        self.config.write(fp)
-        fp.close()
+        self.save_and_reload()
 
     def update_user_info(self, user_info):
         self.config.set(user_info['username'], 'password',
                         user_info['password'])
         self.config.set(user_info['username'], 'ethernet_interface',
                         user_info['ethernet_interface'])
-        fp = open(self.users_cfg_path, 'w')
-        self.config.write(fp)
-        fp.close()
+        self.config.set(user_info['username'], 'dhcp_command',
+                        user_info['dhcp_command'])
+        self.config.set(user_info['username'], 'daemon',
+                        user_info['daemon'])
+        self.save_and_reload()

yah3c/yah3c.py

@@ -4,29 +4,34 @@
 
 """
 
-__version__ = '0.4'
+__version__ = '0.5'
 
 import os, sys
 import ConfigParser
 import getpass
-from socket import socket
+import argparse
+import logging
 
 import eapauth
 import usermgr
-import argparse
+
 
 def parse_arguments():
-    parser = argparse.ArgumentParser(description='Yet Another H3C.', prog='yah3c')
+    parser = argparse.ArgumentParser(description='Yet Another H3C Authentication Client', prog='yah3c')
     parser.add_argument('-u', '--username',
-            help='auth username')
-    parser.add_argument('-p', '--password',
-            help='auth password')
-    parser.add_argument('-i', '--interface', default='eth0',
-            help='Etherent interface used to send packet.eth0 by default.')
-    parser.add_argument('-d', '--daemon', action='store_true',
-            help='Go to background after authentication.')
+            help='Login in with this username')
+    # parser.add_argument('-p', '--password',
+    #         help='Password')
+    # parser.add_argument('-i', '--interface', default='eth0',
+    #         help='Etherent interface used. Set as eth0 by default.')
+    # parser.add_argument('-d', '--daemon', action='store_true',
+    #         help='Fork to background after authentication.')
+    # parser.add_argument('-D', '--dhcp',
+    #         help='DHCP cmd used to obtain ip after authentication.')
+    parser.add_argument('-debug', action='store_true',
+            help='Enable debugging mode')
     args = parser.parse_args()
-    print args
+    return args
 
 def prompt_user_info():
     username = raw_input('Input username: ')
@@ -37,13 +42,26 @@ def prompt_user_info():
             break
         else:
             print 'Password do not match!'
+
     dev = raw_input('Decice(eth0 by default): ')
     if not dev:
         dev = 'eth0'
+
+    choice = raw_input('Forked to background after authentication(Yes by default)\n<Y/N>: ')
+    if choice == 'n' or choice == 'N':
+        daemon = False
+    else:
+        daemon = True
+
+    dhcp_cmd = raw_input('Dhcp command(Press Enter to pass): ')
+    if not dhcp_cmd:
+        dhcp_cmd = ''
     return {
         'username': username,
         'password': password,
-        'ethernet_interface': dev
+        'ethernet_interface': dev,
+        'daemon': daemon,
+        'dhcp_command': dhcp_cmd
     }
 
 def enter_interactive_usermanager():
@@ -80,22 +98,38 @@ def enter_interactive_usermanager():
     else: 
         return users_info[choice - 1]
 
+def start_yah3c(login_info):
+    yah3c = eapauth.EAPAuth(login_info)
+    yah3c.serve_forever()
+
 def main():
-    # TODO: combine cli args with config
     args = parse_arguments()
+    args = vars(args)
 
     # check for root privilege
     if not (os.getuid() == 0):
         print (u'亲，要加sudo!')
         exit(-1)
 
-    if len(sys.argv) == 1:
-        # enter interactive mode
-        login_info = enter_interactive_usermanager()
+    # check if debugging mode enabled
+    if args['debug'] is True:
+        logging.basicConfig(level=logging.DEBUG,
+                format='%(asctime)s %(levelname)s: %(message)s',
+                datefmt='%Y-%m-%d %H:%M:%S')
+        logging.debug('Debugging mode enabled.')
+        logging.debug(args)
 
-    yah3c = eapauth.EAPAuth(login_info)
-    yah3c.serve_forever()
+    # if no username specified then enter interactive mode
+    if args['username'] is None:
+        login_info = enter_interactive_usermanager()
+        logging.debug(login_info)
+        start_yah3c(login_info)
 
+    # if there is username, then get it's info
+    um = usermgr.UserMgr()
+    login_info = um.get_user_info(args['username'])
+    logging.debug(login_info)
+    start_yah3c(login_info)
 
 if __name__ == "__main__":
     main()