Commit 6f28f01b authored by Jason Frisvold's avatar Jason Frisvold
Browse files

- use click for CLI parameters

- rearrange code
- rename config files (to be replaced by config object later)
parent 42206e14
#!/usr/bin/python
#!/usr/bin/env python
# * Skynet - Automated "Cloud" Security Scanner * #
# * Copyright (C) 2014-present Jason Frisvold <friz@godshell.com> * #
......@@ -20,7 +20,7 @@
# Import libraries
import sys
import getopt
import click
# Try/Except to handle Python 3 renaming of configparser
try:
import configparser
......@@ -37,90 +37,20 @@ import re
import shlex
import tempfile
sys.path.append('../libs')
from skynettimer import SkynetTimer
from funcs import pid_exists
from libs.skynettimer import SkynetTimer
from libs.funcs import pid_exists
# Global Variables
verbose = False
programname = 'SkyNet Spawner'
version = '1.0'
configfile = 'spawner.conf'
defaultconfigfile = 'spawner.default.conf'
configfile = 'client.conf'
defaultconfigfile = 'client.conf.default'
cfg = {}
# Set up the logger
logging.basicConfig()
logger = logging.getLogger('spawner')
###
# Main
###
def main(argv):
try:
opts, args = getopt.gnu_getopt(argv,
"dhvVs:c:",
["debug", "help", "verbose", "config",
"version", "license"])
except getopt.GetoptError as err:
print(str(err))
usage()
sys.exit(2)
for o, a in opts:
if o in ("-v", "--verbose"):
logger.setLevel(logging.INFO)
elif o in ("-d", "--debug"):
logger.setLevel(logging.DEBUG)
elif o in ("-c", "--config"):
global configfile
# TODO : This should be validated
configfile = o
elif o in ("-V", "--version"):
print('{0} {1}'.format(programname, version))
sys.exit()
elif o in ("-h", "--help"):
usage()
sys.exit()
elif o in ("--license"):
license()
sys.exit()
else:
assert False, "unhandled option"
loadconfig()
# Main loop
while True:
loopstarttime = datetime.now()
logger.info('Loop starts at {0}'.format(str(loopstarttime)))
timingdb = open_timing_database()
filelist = check_configdir()
if (len(filelist) > 0):
process_config(timingdb, filelist)
spawnlist = check_timing(loopstarttime, timingdb)
spawnlist = check_complete(spawnlist, timingdb)
if (len(spawnlist) > 0):
spawn_process(spawnlist)
# Check for child exit status so we don't have zombies
try:
os.waitpid(-1, os.WNOHANG)
except OSError:
pass
logger.info('Sleeping until next loop iteration')
# Sleep until the beginning of the next minute
sleeptime = 60 - datetime.utcnow().second
time.sleep(sleeptime)
logger = logging.getLogger('client')
def loadconfig():
......@@ -440,27 +370,13 @@ def open_timing_database():
sys.exit()
###
# Usage
###
def usage():
print(('Usage: ' + sys.argv[0] + ' [OPTION]... '))
print('{0}'.format(programname))
print()
print('Mandatory arguments to long options are mandatory for short '
'options too.')
print(' -c <file> configuration file')
print(' -d, --debug debug')
print(' -h, --help display this help and exit')
print(' -v, --verbose verbose output')
print(' -V, --version output version information and exit')
print(' --license output license information and exit')
###
# License
###
def license():
def print_license(ctx, param, value):
if not value or ctx.resilient_parsing:
return
print('Copyright (C) 2014 Jason Frisvold <friz@godshell.com>')
print('License GPLv2+: GNU GPL version 2 or later '
'<http://gnu.org/licenses/gpl.html>.')
......@@ -468,12 +384,69 @@ def license():
print('There is NO WARRANTY, to the extent permitted by law.')
print('')
print('Written by Jason Frisvold.')
ctx.exit()
###
# Main
###
@click.command()
@click.version_option(version=version, prog_name=programname)
@click.option('--license', is_flag=True, callback=print_license,
expose_value=False, is_eager=True, help='Show the license and exit.')
@click.option('--verbose', '-v', count=True, help='Increase verbosity')
@click.option('--config', '-c', help='Specify location of the config file')
def main(verbose, config):
if verbose == 1:
logger.setLevel(logging.INFO)
elif verbose == 2:
logger.setLevel(logging.DEBUG)
elif verbose == 3:
logger.setLevel(logging.DEBUG)
logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
elif verbose == 4:
logger.setLevel(logging.DEBUG)
logging.getLogger('sqlalchemy.engine').setLevel(logging.DEBUG)
loadconfig()
# Main loop
while True:
loopstarttime = datetime.now()
logger.info('Loop starts at {0}'.format(str(loopstarttime)))
timingdb = open_timing_database()
filelist = check_configdir()
if (len(filelist) > 0):
process_config(timingdb, filelist)
spawnlist = check_timing(loopstarttime, timingdb)
spawnlist = check_complete(spawnlist, timingdb)
if (len(spawnlist) > 0):
spawn_process(spawnlist)
# Check for child exit status so we don't have zombies
try:
os.waitpid(-1, os.WNOHANG)
except OSError:
pass
logger.info('Sleeping until next loop iteration')
# Sleep until the beginning of the next minute
sleeptime = 60 - datetime.utcnow().second
time.sleep(sleeptime)
###
# Run the main program, pass in the CLI arguments (minus the command name)
try:
if __name__ == "__main__":
main(sys.argv[1:])
if __name__ == '__main__':
# pylint: disable=no-value-for-parameter
main()
except KeyboardInterrupt:
sys.exit()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment