#!/usr/bin/env python

import argparse
import json
from arachniclient import Client
from arachniclient import Scan

parser = argparse.ArgumentParser(description='Schedule an arachni scan using an arachni REST server.')
parser.add_argument('--url', type=str, nargs=1, help='URL to scan')
parser.add_argument('--hostname', type=str, nargs='?', help='hostname of the arachni REST server (default: 127.0.0.1)')
parser.add_argument('--port', type=int, nargs='?', help='port the arachni REST server listens on (default: 7331)')
parser.add_argument('--delete', type=str, nargs='?', help='delete/abandon a scan on the arachni REST server')
parser.add_argument('--report', type=str, nargs='?', help='download report from the arachni REST server')
parser.add_argument('--format', type=str, nargs='?', help='downloaded report format')
parser.add_argument('--fetch', action="store_true", help='fetch scan ids from the arachni REST server')
parser.add_argument('--status', type=str, nargs='?', help='get status of scan from the arachni REST server')
parser.add_argument('--pause', type=str, nargs='?', help='pause scan running on the arachni REST server')
parser.add_argument('--resume', type=str, nargs='?', help='resume scan running on the arachni REST server')
parser.add_argument('--cookies', type=str, nargs='?', help='cookies to include with the scan')
parser.add_argument('--optionsfile', type=str, nargs='?', help='file locaiton representing the scan options to be used')

args = parser.parse_args()

# Build out the client to use.
client = Client()
if args.hostname and args.port:
    client = Client(hostname=args.hostname, port=args.port)
elif args.hostname:
    client = Client(hostname=args.hostname)
elif args.port:
    client = Client(port=args.port)

# Process the CLI requests
if args.report:
    scan = Scan(client=client, scanId=args.report)
    if args.format:
        print ("trying format: {0}".format(args.format))
    scan.downloadReport(fmt=args.format)

if args.url or args.optionsfile:
    scanOptions = {}
    if args.optionsfile:
        print("Loading json...")
        with open(args.optionsfile, 'r') as f:
            scanOptions = json.loads(f.read())
            
    if args.url:
        url = args.url[0]
        scanOptions.update({"url": url})

    scan = Scan(client=client, scanOptions=scanOptions)
    scan.startScan()
    print("New scan id: {0}".format(scan.id))

if args.fetch:
    print ("Fetching...")
    scans = client.getScans()
    for s in scans:
        print("{0} : {1}".format(s.id, s.scan['status']))

if args.status:
    print("Grabbing scan status...")
    print(client.getScan(args.status))

if args.delete:
    print("Deleting scan {}...".format(args.delete))
    client.deleteScan(args.delete)

if args.pause:
    print("Pausing scan {}...".format(args.pause))
    client.pauseScan(args.pause)

if args.resume:
    print("Resuming scan {}...".format(args.resume))
    client.resumeScan(args.resume)
