# Gitlab bro_hunter's Repository
#
# Original Author: Aaron Eppert (aeppert@gmail.com / aaron@eppert.co )
#

module FILE_EXTRACTION;

type MimeVal: record {
    extension: string;
    protocols: set[string];
};

type MimeConfVal: record {
    mimetype: string;
	extension: string;
    protocols: string;
};

redef FileExtract::prefix = "/data/bro/file_extraction/";

export {
    global _mime_map: table[string] of MimeVal = table();
    global fe_config_filename = @DIR + "./fe.conf" &redef;

     ## Enable MD5 hash for known file types (Default - T)
    const enable_md5_hash_for_known_types = T &redef;

    ## Enable SHA1 hash for known file types (Default - T)
    const enable_sha1_hash_for_known_types = T &redef;

    ## Enable SHA256 hash for known file types (Default - T)
    const enable_sha256_hash_for_known_types = T &redef;
}

function check_protos(proto: string, protocols: set[string]): bool {
    if( (|protocols| > 0 && |proto| > 0) && (proto in protocols || "*" in protocols)) {
        return T;
    }
    return F;
}

# Return based on mt OR if "*" exists as a key, return that as that is the
# catch all case.
function get_MimeVal(mt: string): MimeVal {
    local ret: MimeVal;
    if( |mt| > 0 && (mt in _mime_map)) {
        ret = _mime_map[mt];
    } else {
        if( "*" in _mime_map ) {
            ret = _mime_map["*"];
        }
    }
    return ret;
}

event line(description: Input::EventDescription, tpe: Input::Event, mcv: MimeConfVal) {
	local tmp: MimeVal;

    if ( mcv?$extension ) {
    	tmp$extension = mcv$extension;

        if ( mcv?$protocols ) {
        	local tmp_sv = split_string(mcv$protocols, / /);

        	for(proto in tmp_sv) {
        	 	add tmp$protocols[tmp_sv[proto]];
        	}
        }

    	_mime_map[mcv$mimetype] = tmp;
    }
}

event zeek_init() &priority=20 {
	Input::add_event([$source=fe_config_filename, $name="feconf", $fields=MimeConfVal, $ev=line, $mode=Input::REREAD]);
}

event file_sniff(f: fa_file, meta: fa_metadata) {
    local ext = "";

	if( meta?$mime_type ) {
        local mv = get_MimeVal(meta$mime_type);
        if ( mv?$extension ) {
            ext = mv$extension;

            if ( !check_protos(f$source, mv$protocols) ) {
    		    return;
    	    }

            local output_dir = fmt("%s", FileExtract::prefix);
            local fe_dir: vector of string = vector( f$source,
                                             fmt("%s", strftime("%Y", current_time())),
                                             fmt("%s", strftime("%m", current_time())),
                                             fmt("%s", strftime("%d", current_time())),
                                             fmt("%s", strftime("%H", current_time())),
                                             fmt("%s", strftime("%M", current_time())));

            for( d in fe_dir ) {
                output_dir = output_dir + "/" + fe_dir[d];
                mkdir(output_dir);
            }

            local time_dir = fmt("%s", strftime("%Y/%m/%d/%H/%M", current_time()));
            local fname = fmt("%s/%s/%s", f$source, time_dir, f$id);

            if ( |ext| != 0 ) {
                fname = fmt("%s/%s/%s.%s", f$source, time_dir, f$id, ext);
        
                if ( FILE_EXTRACTION::enable_md5_hash_for_known_types ) {
                    Files::add_analyzer(f, Files::ANALYZER_MD5);
                }

                if ( FILE_EXTRACTION::enable_sha1_hash_for_known_types ) {
                    Files::add_analyzer(f, Files::ANALYZER_SHA1);
                }

                if ( FILE_EXTRACTION::enable_sha256_hash_for_known_types ) {
                    Files::add_analyzer(f, Files::ANALYZER_SHA256);
                }
                
                Files::add_analyzer(f, Files::ANALYZER_EXTRACT, [$extract_filename=fname]);
            } 
        }
    }
}