Doc fixes

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

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)