diff --git a/dga_detector.py b/dga_detector.py old mode 100644 new mode 100755 index a8fd7bd..2c96979 --- a/dga_detector.py +++ b/dga_detector.py @@ -1,3 +1,4 @@ +#!/usr/bin/env python3 import pickle from gib import gib_detect_train from dga_routines import count_consonants, entropy @@ -7,12 +8,27 @@ def read_file(filename): + """ + Read each line from a file (with newline removed). + + :param filename: file name/path to read from + :type filename: str + :returns: generator of str + """ with open(filename) as f: for line in f: yield line.strip("\n") def domain_check(domain): + """ + Check a given domain. This function operates on the second level domain, e.g. on 'example' of 'example.net'. + + :param domain: domain to check + :type domain: str + :returns: tuple (second level domain - str, entropy - float, number of consonants - int, length - int), may be empty + :note: Domains shorter than six characters, localized domains (i.e. 'xn-') and onion services (i.e. '.onion') are not processed. + """ # skip tor domains if domain.endswith(".onion"): print("Tor domains is ignored...") @@ -33,7 +49,7 @@ def domain_check(domain): return domain_without_sub, domain_entropy, domain_consonants, domain_length -def main(): +if __name__ == "__main__": parser = argparse.ArgumentParser(description="DGA domain detection") parser.add_argument("-d", "--domain", help="Domain to check") parser.add_argument("-f", "--file", help="File with domains. One per line") @@ -93,4 +109,3 @@ def main(): /_____/ \____/ /_/ |_| /_____/ \___/\__/ \___/\___/ \__/ \____//_/ ''') parser.print_help() -main() diff --git a/dga_routines.py b/dga_routines.py index 889605b..264f21b 100755 --- a/dga_routines.py +++ b/dga_routines.py @@ -5,7 +5,11 @@ def entropy(string): """ - Calculates the Shannon entropy of a string + Calculates the Shannon entropy of a string. + + :param string: input string + :type string: str + :returns: float """ # get probability of chars in string @@ -19,7 +23,11 @@ def entropy(string): def count_consonants(string): """ - Counting consonants in a string + Counting consonants in a string. + + :param string: input string + :type string: str + :retuns: int (number of consonants) """ consonants = re.compile("[bcdfghjklmnpqrstvwxyz]") count = consonants.findall(string)