from subprocess import Popen, PIPE
from parsers.tail import read_tailed_files
from db import Database
from util import LOGGER

def parse_args():
    from argparse import ArgumentParser
    parser = ArgumentParser(description='Snapshot statistics from a machine')
    parser.add_argument('--ip', default='',
                        help='connect to a remote host (recommended)')
    # Multiple pids could be set using bash expansion: {1234,2345}
    parser.add_argument('-p', '--pid', default='*',
                        help='the pid(s) to look up (default: *)')
    parser.add_argument('-u', '--user', default='root',
                        help='user to log into remote host with (default: root)')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='enable more verbose logging')
    parser.add_argument('--overwrite', action='store_true',
                        help='always overwrite the db, even if it exists')
    parser.add_argument('--db', required=True,
                        help='path to store the data to (sqlite format)')
    args = parser.parse_args()
    return args


def read_stats(args):
    # This is the command to grap all of the necessary info.
    # Note that -v is passed to tail - this is so we always the filename
    # given to us, which is needed for parsing.
    # As processes can be transient, we can get errors here about
    # non-existant files, so ignore them, this is expectedself.
    cmd = 'bash -c "'\
          'tail -v -n +1 '\
              '/proc/%s/{cmdline,smaps} '\
              '/proc/meminfo '\
              '/proc/loadavg '\
              '/proc/uptime '\
          '2>/dev/null;' \
          'tail -v -n +1 '\
              '/proc/%s/stat '\
          '2>/dev/null | '\
          'awk \''\
              '/==>/ {print \$0} '\
              '/^[0-9]/ {print \$2, \$10, \$12, \$14, \$15, \$22}\';"'


    # Accept a space-separated list of pids as that is what pidof(8) returns and
    # it's quite likely you'll want to invoke this script with something like:
    #
    #     --pid "`pidof foobar`"
    #
    # at some point.
    if args.pid.isdigit() or args.pid == '*':
        pids = args.pid
    else:
        pids = '{%s}' % args.pid.replace(' ', ',')

    if args.ip == '':
        LOGGER.info('Loading local procfs files')
        cmd = "sudo %s" % (cmd % (pids, pids))
        stream = Popen(cmd, shell=True, bufsize=-1, stdout=PIPE).stdout
    elif args.ip != '':
        cmd = """ssh %s@%s '%s'""" % (args.user, args.ip, cmd % (pids, pids))
        stream = Popen(cmd, shell=True, bufsize=-1, stdout=PIPE).stdout

    LOGGER.info('Reading procfs with cmd: %s' % cmd)
    return read_tailed_files(stream)


def main(args):
    import logging
    if args.verbose:
        LOGGER.setLevel(logging.DEBUG)
    else:
        LOGGER.setLevel(logging.INFO)

    # Get the database handle
    db = Database(args.db, args.overwrite)
    # Read all the data we need
    system_stats, processes, memory_stats = read_stats(args)

    LOGGER.info('Found {} process(es) and {} used memory fragments'.format(
                len(processes), len(memory_stats)))
    LOGGER.info('Regions: %s' % memory_stats)

    db.add(args.ip if len(args.ip) else '[local]', system_stats, memory_stats, processes)

if __name__ == '__main__':
    main(parse_args())
@@ -28,16 +28,28 @@
 n, PIPE%0A
+import time%0A
 from par
@@ -1095,34 +1095,476 @@
     
-args = parser.parse_args()
+parser.add_argument('-c', '--count', default=1, type=int,%0A                        help='the number of snapshots to collect')%0A    parser.add_argument('--period', type=int, default=0,%0A                        help='number of seconds between subsequent snapshots')%0A    args = parser.parse_args()%0A%0A    if args.count %3E 1 and 0 == args.period:%0A        print ('Error: You must set the period if count %3E 1%5Cn')%0A        parser.print_help()%0A        import sys%0A        sys.exit(1)%0A
 %0A   
@@ -3378,16 +3378,107 @@
 rwrite)%0A
+%0A    for i in range(args.count):%0A        if i %3E 0:%0A            time.sleep(args.period)%0A    
     # Re
@@ -3501,16 +3501,20 @@
 we need%0A
+    
     syst
@@ -3567,24 +3567,28 @@
 (args)%0A%0A    
+   
+ 
 LOGGER.info(
@@ -3663,16 +3663,20 @@
         
+    
 len(proc
@@ -3703,24 +3703,28 @@
 tats)))%0A    
+  
+  
 LOGGER.info(
@@ -3750,24 +3750,28 @@
 ory_stats)%0A%0A
+    
     db.add(a

Add periodic snapshotting.