#! /usr/bin/env python

# This file is part of IVRE.
# Copyright 2011 - 2014 Pierre LALET <pierre.lalet@cea.fr>
#
# IVRE is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# IVRE is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
# License for more details.
#
# You should have received a copy of the GNU General Public License
# along with IVRE. If not, see <http://www.gnu.org/licenses/>.

import ivre.db

import os

def recursive_filelisting(base_directories):
    "Iterator on filenames in base_directories"

    for base_directory in base_directories:
        for root, _, files in os.walk(base_directory):
            for leaffile in files:
                yield os.path.join(root, leaffile)

def main():
    try:
        import argparse
        parser = argparse.ArgumentParser(
            description='Parse NMAP scan results and add them in DB.')
        parser.add_argument('scan', nargs='*', metavar='SCAN',
                            help='Scan results')

    except ImportError:
        import optparse
        parser = optparse.OptionParser(
            description='Parse NMAP scan results and add them in DB.')
        parser.parse_args_orig = parser.parse_args

        def my_parse_args():
            res = parser.parse_args_orig()
            res[0].ensure_value('scan', res[1])
            return res[0]
        parser.parse_args = my_parse_args
        parser.add_argument = parser.add_option

    parser.add_argument('-c', '--categories', default='',
                        help='Scan categories.')
    parser.add_argument('-s', '--source', default=None,
                        help='Scan source.')
    parser.add_argument('-t', '--test', action='store_true',
                        help='Test mode.')
    parser.add_argument('--port', action='store_true',
                        help='Need ports.')
    parser.add_argument('--never-archive', action='store_true',
                        help='Never archive.')
    parser.add_argument('--archive', '--archive-same-host', action='store_true',
                        help='Archive results for the same host.')
    parser.add_argument('--archive-same-host-and-source', action='store_true',
                        help='Archive results with both the same host and'
                        ' source (this is the default).')
    parser.add_argument('-r', '--recursive', action='store_true',
                        help='Import all files from given directories.')
    args = parser.parse_args()
    categories = []
    source = None
    database = ivre.db.db.nmap
    categories = args.categories.split(',')
    if args.test:
        database = ivre.db.DBNmap()
    if args.never_archive:
        def gettoarchive(collection, addr, source):
            return []
    elif args.archive:
        def gettoarchive(collection, addr, source):
            return collection.find({'addr': addr})
    else:  #args.archive_same_host_and_source
        def gettoarchive(collection, addr, source):
            return collection.find({'addr': addr,
                                    'source': source})
    if args.recursive:
        scans = recursive_filelisting(args.scan)
    else:
        scans = args.scan
    count = 0
    for scan in scans:
        if database.store_scan(
                scan,
                categories=categories, source=args.source,
                needports=args.port, gettoarchive=gettoarchive
        ):
            count += 1
    print "%d results imported." % count

if __name__ == '__main__':
    main()
