| #!/usr/bin/env python3 |
| # -*- coding: utf-8 -*- |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| #the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import os |
| import sys |
| import threading |
| import multiprocessing |
| import yaml |
| import json |
| import time |
| import re |
| import argparse |
| import plugins.scanners |
| import plugins.brokers.kibbleES |
| #import plugins.kibbleJSON |
| |
| VERSION = "0.1.0" |
| CONFIG_FILE = "conf/config.yaml" |
| PENDING_OBJECTS = [] |
| BIG_LOCK = threading.Lock() |
| |
| def base_parser(): |
| arg_parser = argparse.ArgumentParser() |
| arg_parser.add_argument("-o", "--org", help="The organisation to gather stats for. If left out, all organisations will be scanned.") |
| arg_parser.add_argument("-f", "--config", help="Location of the yaml config file (full path)") |
| arg_parser.add_argument("-a", "--age", help="Minimum age in hours before performing a new scan on an already processed source. --age 12 will not process any source that was processed less than 12 hours ago, but will process new sources.") |
| arg_parser.add_argument("-s", "--source", help="A specific source (wildcard) to run scans on.") |
| arg_parser.add_argument("-n", "--nodes", help="Number of nodes in the cluster (used for load balancing)") |
| arg_parser.add_argument("-t", "--type", help="Specific type of scanner to run (default is run all scanners)") |
| arg_parser.add_argument("-e", "--exclude", nargs = '+', help="Specific type of scanner(s) to exclude") |
| arg_parser.add_argument("-v", "--view", help="Specific source view to scan (default is scan all sources)") |
| return arg_parser |
| |
| def pprint(string, err = False): |
| line = "[core]: %s" % (string) |
| if err: |
| sys.stderr.write(line + "\n") |
| else: |
| print(line) |
| |
| |
| def isMine(ID, config): |
| if config['scanner'].get('balance', None): |
| a = config['scanner']['balance'].split('/') |
| nodeNo = int(a[0]) |
| numNodes = int(a[1]) |
| if numNodes == 0: |
| return True |
| bignum = int(ID, 16) % numNodes |
| if bignum == int(nodeNo)-1: |
| return True |
| return False |
| return True |
| |
| class scanThread(threading.Thread): |
| """ A thread object that grabs an item from the queue and processes |
| it, using whatever plugins will come out to play. """ |
| def __init__(self, broker, org, i, t = None, e = None): |
| super(scanThread, self).__init__() |
| self.broker = broker |
| self.org = org |
| self.id = i |
| self.bit = self.broker.bitClass(self.broker, self.org, i) |
| self.stype = t |
| self.exclude = e |
| pprint("Initialized thread %i" % i) |
| |
| def run(self): |
| global BIG_LOCK, PENDING_OBJECTS |
| time.sleep(0.5) # Primarily to align printouts. |
| # While there are objects to snag |
| a = 0 |
| while PENDING_OBJECTS: |
| BIG_LOCK.acquire(blocking = True) |
| try: |
| # Try grabbing an object (might not be any left!) |
| obj = PENDING_OBJECTS.pop(0) |
| except: |
| pass |
| BIG_LOCK.release() |
| if obj: |
| # If load balancing jobs, make sure this one is ours |
| if isMine(obj['sourceID'], self.broker.config): |
| # Run through list of scanners in order, apply when useful |
| for sid, scanner in plugins.scanners.enumerate(): |
| |
| if scanner.accepts(obj): |
| self.bit.pluginname = "plugins/scanners/" + sid |
| # Excluded scanner type? |
| if self.exclude and sid in self.exclude: |
| continue |
| # Specific scanner type or no types mentioned? |
| if not self.stype or self.stype == sid: |
| scanner.scan(self.bit, obj) |
| else: |
| break |
| self.bit.pluginname = "core" |
| self.bit.pprint("No more objects, exiting!") |
| |
| def main(): |
| pprint("Kibble Scanner v/%s starting" % VERSION) |
| global CONFIG_FILE, PENDING_OBJECTS |
| args = base_parser().parse_args() |
| |
| # Load config yaml |
| if args.config: |
| CONFIG_FILE = args.config |
| config = yaml.load(open(CONFIG_FILE)) |
| pprint("Loaded YAML config from %s" % CONFIG_FILE) |
| |
| # Which broker type do we use here? |
| broker = None |
| if 'elasticsearch' in config and config['elasticsearch'].get('enabled', False): |
| pprint("Using direct ElasticSearch broker model") |
| broker = plugins.brokers.kibbleES.Broker(config) |
| else: |
| pprint("Using HTTP JSON broker model") |
| broker = plugins.brokers.kibbleJSON.Broker(config) |
| |
| orgNo = 0 |
| sourceNo = 0 |
| for org in broker.organisations(): |
| if not args.org or args.org == org.id: |
| pprint("Processing organisation %s" % org.id) |
| orgNo += 1 |
| |
| # Compile source list |
| # If --age is passed, only append source that either |
| # have never been scanned, or have been scanned more than |
| # N hours ago by any scanner. |
| if args.age: |
| minAge = time.time() - int(args.age) * 3600 |
| for source in org.sources(view=args.view): |
| tooNew = False |
| if 'steps' in source: |
| for key, step in source['steps'].items(): |
| if 'time' in step and step['time'] >= minAge: |
| tooNew = True |
| break |
| if not tooNew: |
| if not args.source or (args.source == source['sourceID']): |
| PENDING_OBJECTS.append(source) |
| else: |
| PENDING_OBJECTS = [] |
| for source in org.sources(view=args.view): |
| if not args.source or (args.source == source['sourceID']): |
| PENDING_OBJECTS.append(source) |
| sourceNo += len(PENDING_OBJECTS) |
| |
| # Start up some threads equal to number of cores on the box, |
| # but no more than 4. We don't want an IOWait nightmare. |
| threads = [] |
| core_count = min((4, int( multiprocessing.cpu_count() ))) |
| for i in range(0, core_count): |
| sThread = scanThread(broker, org, i+1, args.type, args.exclude) |
| sThread.start() |
| threads.append(sThread) |
| |
| # Wait for them all to finish. |
| for t in threads: |
| t.join() |
| |
| pprint("All done scanning for now, found %i organisations and %i sources to process." % (orgNo, sourceNo)) |
| |
| if __name__ == '__main__': |
| main() |
