18a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen#!/usr/bin/env python
28a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
38a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen"""This spawns a sub-shell (bash) and gives the user interactive control. The
48a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenentire shell session is logged to a file called script.log. This behaves much
58a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenlike the classic BSD command 'script'.
68a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
78a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen./script.py [-a] [-c command] {logfilename}
88a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
98a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    logfilename : This is the name of the log file. Default is script.log.
108a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    -a : Append to log file. Default is to overwrite log file.
118a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    -c : spawn command. Default is to spawn the sh shell.
128a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
138a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny ChenExample:
148a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
158a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    This will start a bash shell and append to the log named my_session.log:
168a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
178a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        ./script.py -a -c bash my_session.log
188a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
198a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen"""
208a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
218a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenimport os, sys, time, getopt
228a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenimport signal, fcntl, termios, struct
238a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenimport traceback
248a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenimport pexpect
258a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
268a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenglobal_pexpect_instance = None # Used by signal handler
278a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
288a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chendef exit_with_usage():
298a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
308a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    print globals()['__doc__']
318a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    os._exit(1)
328a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
338a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chendef main():
348a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
358a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    ######################################################################
368a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    # Parse the options, arguments, get ready, etc.
378a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    ######################################################################
388a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    try:
398a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        optlist, args = getopt.getopt(sys.argv[1:], 'h?ac:', ['help','h','?'])
408a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    except Exception, e:
418a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        print str(e)
428a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        exit_with_usage()
438a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    options = dict(optlist)
448a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    if len(args) > 1:
458a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        exit_with_usage()
468a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
478a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    if [elem for elem in options if elem in ['-h','--h','-?','--?','--help']]:
488a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        print "Help:"
498a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        exit_with_usage()
508a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
518a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    if len(args) == 1:
528a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        script_filename = args[0]
538a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    else:
548a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        script_filename = "script.log"
558a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    if '-a' in options:
568a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        fout = file (script_filename, "ab")
578a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    else:
588a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        fout = file (script_filename, "wb")
598a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    if '-c' in options:
608a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        command = options['-c']
618a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    else:
628a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        command = "sh"
638a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
648a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    # Begin log with date/time in the form CCCCyymm.hhmmss
658a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    fout.write ('# %4d%02d%02d.%02d%02d%02d \n' % time.localtime()[:-3])
668a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
678a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    ######################################################################
688a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    # Start the interactive session
698a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    ######################################################################
708a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    p = pexpect.spawn(command)
718a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    p.logfile = fout
728a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    global global_pexpect_instance
738a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    global_pexpect_instance = p
748a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    signal.signal(signal.SIGWINCH, sigwinch_passthrough)
758a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
768a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    print "Script recording started. Type ^] (ASCII 29) to escape from the script shell."
778a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    p.interact(chr(29))
788a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    fout.close()
798a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    return 0
808a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
818a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chendef sigwinch_passthrough (sig, data):
828a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
838a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    # Check for buggy platforms (see pexpect.setwinsize()).
848a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    if 'TIOCGWINSZ' in dir(termios):
858a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        TIOCGWINSZ = termios.TIOCGWINSZ
868a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    else:
878a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        TIOCGWINSZ = 1074295912 # assume
888a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    s = struct.pack ("HHHH", 0, 0, 0, 0)
898a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    a = struct.unpack ('HHHH', fcntl.ioctl(sys.stdout.fileno(), TIOCGWINSZ , s))
908a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    global global_pexpect_instance
918a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    global_pexpect_instance.setwinsize(a[0],a[1])
928a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
938a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chenif __name__ == "__main__":
948a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    try:
958a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        main()
968a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    except SystemExit, e:
978a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        raise e
988a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen    except Exception, e:
998a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        print "ERROR"
1008a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        print str(e)
1018a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        traceback.print_exc()
1028a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen        os._exit(1)
1038a3c0430323c28c1fbe8ceecd2cd8e58b64a9295Johnny Chen
104