#!/usr/bin/env python
import pyclamd
import boto3
import pprint
import sys
import argparse
import logging as log
import os
import threading
from scanner import Scanner

sema = threading.Semaphore(value=9) # Boto3 has a limit of 10 parallel connections
threads = list()


def _run(bucketname, object_key):
    sema.acquire()

    scanner = Scanner(args.clamd_host, args.clamd_port, args.tagging, {'clean': args.tag_clean, 'infected': args.tag_infected})
    scanner.file_check(bucketname, object_key, args.tagging)

    sema.release()

def _fullscan(bucketname):

    s3 = boto3.client("s3")
    all_objects = s3.list_objects(Bucket=bucketname) 

    for s3file in all_objects['Contents']:
        x = threading.Thread(target=_run, args=(bucketname, s3file['Key']))
        threads.append(x)
        x.start()

    for y in threads:
        y.join()

    print('DONE')

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose", action='store_true', required=False, help="increase output verbosity")
    parser.add_argument("-vv", "--veryverbose", action='store_true', required=False, help="increase output verbosity")
    parser.add_argument("-b", "--aws_bucket", default=os.environ.get('AWS_BUCKET'), required=False, help="Bucket name")
    parser.add_argument("-c", "--clamd_host", default=os.environ.get('CLAMD_HOST', '127.0.0.1'), required=False, help="Host or IP of clamd Server")
    parser.add_argument("-p", "--clamd_port", default=os.environ.get('CLAMD_PORT', '3310'), required=False, help="Port of clamd Server")
    parser.add_argument("-t", "--tagging", default=os.environ.get('TAGGING', 'False'), required=False, help="Enable tagging")
    parser.add_argument("-tc", "--tag_clean", default=os.environ.get('TAG_CLEAN', 'clean'), required=False, help="Tag to be set when CLEAN")
    parser.add_argument("-ti", "--tag_infected", default=os.environ.get('TAG_INFECTED', 'infected'), required=False, help="Tag to be set when INFECTD")

    args = parser.parse_args()

    if args.verbose:
        log.basicConfig(format='%(levelname)s:%(message)s', level=log.INFO)

    if args.veryverbose:
        log.basicConfig(format='%(levelname)s:%(message)s', level=log.DEBUG)


    log.info("Clamd Host: {}:{}".format(args.clamd_host, args.clamd_port))

    try:
        _fullscan(args.aws_bucket)
    except KeyboardInterrupt:
        print('EXIT: User aborted scan')
    
    sys.exit()