Shannon-Entropy.py

import sys, os
import argparse
import math
import pefile

'''
Shannon-Entropy.py
Determine Shannon Entropy of any file - value output between 0 (ordered) and 8 (fully random)
Optinally add `-pe` argument to parse file headers of PE file and determine entropy for each header
'''

def main(args=sys.argv[1:]):
    parser = argparse.ArgumentParser(description="Calculate Shannon entropy of input file")
    parser.add_argument('-file', '-f', help="Input file")
    parser.add_argument('-pe', help="BOOL: Specify if input file is PE for additional analysis", action="store_true")
    args = parser.parse_args(args)
    inFile = args.file
    pe = args.pe

    if not inFile:
        print("[-] ERROR! Missing input file parameter '-f'")
        sys.exit()

    
    file_path = inFile

    if not os.path.exists(file_path):
        print('File not found:', file_path)
        sys.exit()
    
    # Calculate entropy of input file
    entropy = calc_entropy_of_file(file_path)
    print("File Shannon Entropy:", entropy)

    # Do further analysis on headers if input file is PE
    if pe:
        peFile = pefile.PE(file_path, fast_load=True)
        # Enumerate each section to determine entropy
        for section in peFile.sections:
            print(section.Name.decode("utf-8"))
            print("Section entropy:", str(shannon_entropy(section.get_data())))



def shannon_entropy(data):
    # Determine the frequency of each byte value
    byte_counts = [0] * 256
    for byte in data:
        byte_counts[byte] += 1

    # Determine the probability of each byte value
    total_bytes = len(data)
    probabilities = [count / total_bytes for count in byte_counts if count > 0]

    # Determine Shannon entropy
    entropy = -sum(p * math.log2(p) for p in probabilities)
    return entropy


def calc_entropy_of_file(file_path):
    with open(file_path, 'rb') as file:
        data = file.read()
        entropy = shannon_entropy(data)
        return entropy

if __name__ == '__main__':
    main()