#!/usr/bin/env python3

# Copyright 2014-2023 Security Onion Solutions, LLC
#
# This program 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.
#
# This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.

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)

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")
