#!/usr/bin/env python3

# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one
# or more contributor license agreements. Licensed under the Elastic License 2.0 as shown at
# https://securityonion.net/license; you may not use this file except in compliance with the
# Elastic License 2.0.

import os
import shutil
import time
import hashlib
import logging
import yaml
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

with open("/opt/so/conf/strelka/filecheck.yaml", "r") as ymlfile:
    cfg = yaml.load(ymlfile, Loader=yaml.Loader)

extract_path = cfg["filecheck"]["extract_path"]
historypath = cfg["filecheck"]["historypath"]
strelkapath = cfg["filecheck"]["strelkapath"]
logfile = cfg["filecheck"]["logfile"]
recycle_secs = cfg["filecheck"].get("recycle_secs", 300)

logging.basicConfig(filename=logfile, filemode='w', format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.INFO)

def checkexisting():
    logging.info("Checking for existing files");
    for root, dirs, files in os.walk(extract_path):
        for file in files:
            try:
                path = os.path.join(root, file)
                filename = os.path.join(extract_path, path)
                checksum(filename) 
            except Exception as err:
                logging.error("Failed to process file: " + file)

def checksum(filename):
    if os.path.isfile(filename) and "/tmp/" not in filename:
        with open(filename, 'rb') as afile:
            logging.info("Processing file: " + filename)
            shawnuff = hashlib.sha1()
            buf = afile.read(8192)
            while len(buf) > 0:
                shawnuff.update(buf)
                buf = afile.read(8192)
            hizash=shawnuff.hexdigest()
            process(filename, hizash)

def process(filename, hizash):
    if os.path.exists(historypath + hizash):
        logging.info(filename + " Already exists.. removing")
        os.remove(filename)
    else:
        # Write the file
        logging.info(filename + " is new. Creating a record and sending to Strelka")
        with open(os.path.join(historypath + hizash), 'w') as fp:
            pass
        head, tail = os.path.split(filename)

        # Move the file
        shutil.move(filename, strelkapath + tail)

class CreatedEventHandler(FileSystemEventHandler):
    def on_created(self, event):
        logging.info("File create detected: " + event.src_path)
        checksum(event.src_path)

    def on_moved(self, event):
        logging.info("File move detected: " + event.src_path + " -> " + event.dest_path)
        checksum(event.dest_path)

if __name__ == "__main__":
    logging.info("Starting filecheck")

    event_handler =CreatedEventHandler()

    shutdown = False
    while not shutdown:
        checkexisting()
        logging.info("Scheduling observer")
        observer = Observer()
        observer.schedule(event_handler, extract_path, recursive=True)
        observer.start()
        try:
            time.sleep(recycle_secs)
        except KeyboardInterrupt:
            logging.warn("User requested shutdown")
            shutdown = True

        observer.stop()
        observer.join()

        if not shutdown:
            logging.info("Recycling observer to pick up new subdirectories")

    logging.info("Exiting filecheck")
