#!/usr/bin/env python

"""
OVERVIEW

Extract USB mass storage device events from Cb Enterprise Response (CbER).
"""

import argparse
import csv
import json
import os
import sys

from cbapi.response import CbEnterpriseResponseAPI
from cbapi.response.models import Process


match_guid = '{53f56307-b6bf-11d0-94f2-00a0c91efb8b}'
search_terms = ["registry\\machine\\system\\currentcontrolset\\control\\deviceclasses\\{53f5630d-b6bf-11d0-94f2-00a0c91efb8b}\\*",
                "registry\\machine\\currentcontrolset\\control\\deviceclasses\\{53f56307-b6bf-11d0-94f2-00a0c91efb8b}\\*"]


class USBEvent:
    def __init__(self, path):
        self.path = path

        self.vendor = ''
        self.product = ''
        self.version = ''
        self.serial = ''
        #self.drive_letter = ''
        #self.volume_name = ''

        self.parse()

    def __repr__(self):
        for k,v in self.__dict__.iteritems():
            print '%s,%s' % (k, v)

    def parse(self):
        path = self.path.split('usbstor#disk&')[1]
        fields = path.split('&')
        self.vendor = fields[0].split('ven_')[1]
        self.product = fields[1].split('prod_')[1]

        if self.vendor == 'drobo':
            # Drobo doesn't provide a version
            drobo_fields = self.product.split('#')
            self.product = drobo_fields[0]
            self.serial = drobo_fields[1]
        else:
            self.version = fields[2].split('#')[0].split('rev_')[1]
            self.serial = fields[2].split('#')[1]


def usbstor_search(cb_conn, query, query_base=None, timestamps=False):

    if query_base is not None:
        query += query_base

    query_result = cb_conn.select(Process).where(query)
    query_result_len = len(query_result)

    results = set()

    for proc in query_result:
        for regmod in proc.regmods:
            if match_guid in regmod.path and 'usbstor#disk&' in regmod.path:
                usb_result = USBEvent(regmod.path)

                output_fields = [proc.hostname,
                                usb_result.vendor,
                                usb_result.product,
                                usb_result.version,
                                usb_result.serial]
                if timestamps == True:
                    output_fields.insert(0, proc.timestamp)

                results.add(tuple(output_fields))

    return results

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--prefix", type=str, action="store",
                        help="Output filename prefix.")
    parser.add_argument("--days", type=int, action="store",
                        help="Number of days to search.")
    parser.add_argument("--minutes", type=int, action="store",
                        help="Number of days to search.")
    parser.add_argument("--timestamps", action="store_true",
                        help="Include timestamps in results.")
    parser.add_argument("--profile", type=str, action="store",
                        help="The credentials.response profile to use.")

    args = parser.parse_args()

    if args.prefix:
        output_filename = '%s-usb-storage-events.csv' % args.prefix
    else:
        output_filename = 'usb-storage-events.csv'

    if args.profile:
        cb = CbEnterpriseResponseAPI(profile=args.profile)
    else:
        cb = CbEnterpriseResponseAPI()

    output_file = file(output_filename, 'w')
    writer = csv.writer(output_file, quoting=csv.QUOTE_ALL)

    header_row = ['endpoint', 'vendor', 'product', 'version', 'serial']
    if args.timestamps == True:
        header_row.insert(0, 'timestamp')
    writer.writerow(header_row)

    for term in search_terms:
        query = 'regmod:%s' % term

        if args.days:
            query += ' start:-%dm' % (args.days*1440)
        elif args.minutes:
            query += ' start:-%dm' % args.minutes

        results = usbstor_search(cb, query, query_base=None, timestamps=args.timestamps)

        for row in results:
            row = list(row)
            row = [col.encode('utf8') if isinstance(col, unicode) else col for col in row]
            writer.writerow(row)

    output_file.close()


if __name__ == '__main__':

    sys.exit(main())