0
\$\begingroup\$

I have always been known to write really sloppy code. Can someone help me rewrite the following code in a neater manner?

import os
import hashlib
import sys
import time

file_list = []

rootdir = "C:/"

print("Program starting!")
print("[+]Collecting virus definitions and allocating memory[+]")

for subdir, dirs, files in os.walk(rootdir):
    for file in files:
        #print os.path.join(subdir, file)
        filepath = subdir + os.sep + file

        if filepath.endswith(".exe") or filepath.endswith(".dll"):
            file_list.append(filepath)
            #print(filepath)

print("[+]Virus definition and memory allocation complete...[+]")
print("[+]Starting scan...[+]")
def countdown():
    for x in range(4):
        print(x+1)
        time.sleep(1)

countdown()

def Scan():
    infected_list = []
    for f in file_list:
        virus_defs = open("VirusLIST.txt", "r")
        file_not_read = False
        print("\nScanning: {}".format(f))
        hasher = hashlib.md5()
        try:
            with open(f, "rb") as file:
                try:
                    buf = file.read()
                    file_not_read = True
                    hasher.update(buf)
                    FILE_HASHED = hasher.hexdigest()
                    print("File md5 checksum: {}".format(FILE_HASHED))
                    for line in virus_defs:
                        if FILE_HASHED == line.strip():
                            print("[!]Malware Detected[!] | File name: {}".format(f))
                            infected_list.append(f)
                        else:
                            pass
                except Exception as e:
                    print("Could not read file | Error: {}".format(e))
        except:
            pass
    print("Infected files found: {}".format(infected_list))
    deleteornot = str(input("Would you like to delete the infected files (y/n): "))
    if deleteornot.upper() == "Y":
        for infected in infected_list:
            os.remove(infected)
            print("File removed: {}".format(infected))
    else:
        print("Executed with exit code 0")
        os.system("PAUSE")
Scan()
\$\endgroup\$
2
  • 7
    \$\begingroup\$ Could you please add an explanation about what the code is supposed to do? \$\endgroup\$
    – Snowbody
    Dec 23, 2017 at 3:50
  • 1
    \$\begingroup\$ Checking executables against hashes is a really outdated (and honestly pretty useless) approach to antivirus software... \$\endgroup\$
    – Vogel612
    Dec 24, 2017 at 12:39

1 Answer 1

10
\$\begingroup\$
rootdir = "C:/"

What operating system is this for? It's not correct for Windows -- Windows is supposed to use the \ character as a path separator. Sometimes you are using os as is recommended, sometimes you roll your own. This is not good.

I think the scan directory ought to be a parameter, not a constant hardcoded in the program. Then put it together with os.path.join(). Why did you not use it?

for subdir, dirs, files in os.walk(rootdir):
     for file in files:
        #print os.path.join(subdir, file)
        filepath = subdir + os.sep + file

        if filepath.endswith(".exe") or filepath.endswith(".dll"):
            file_list.append(filepath)
            #print(filepath)
  • What's the purpose of the commented-out lines?
  • Why do work before you know if the file is one that you need to scan? Only build the temporary string if you need it.
  • This doubly-nested loop could, and probably should, be rewritten as a one line list comprehension.
print("[+]Virus definition and memory allocation complete...[+]")
print("[+]Starting scan...[+]")

def countdown():
     for x in range(4):
         print(x+1)
         time.sleep(1)
countdown()
  • This function is called countdown but it actually counts up. Why?
  • Why is it declared in the middle of imperative coden? Is there any benefit to declaring it and then immediately using it?
  • If you want the numbers from 1 to 4, ask for them. range() can take a second parameter to indicate a range.
def Scan():
  • I note you're using global variables. This is a bad practice. Use function parameters instead.
infected_list = []
    for f in file_list:
        virus_defs = open("VirusLIST.txt", "r")
  • So, you open and read through the entirety of VirusLIST.txt for each file on your hard drive. Does this seem right to you? Perhaps there's a data structure you should be using.
  • Please use a more descriptive name for your loop index variable than f
   file_not_read = False
  • This is confusing. Most people don't deal with double negatives easily. What does this variable actually represent? When does it change from False to True?
   print("\nScanning: {}".format(f))
  • Have you heard about the new formatted string syntax? In this line it would be f"Scanning: {f}"
   hasher = hashlib.md5()
         try:
            with open(f, "rb") as file:
                try:
  • You should just have the outer try:, and catch the appropriate exception rather than catching everything.
   buf = file.read()
     file_not_read = True
     hasher.update(buf)
  • I don't think the temporary variable is necessary
   FILE_HASHED = hasher.hexdigest()
  • Please choose another variable name. This one implies that the variable contains a file; it does not. Also, why is the variable name in all caps?
   print("File md5 checksum: {}".format(FILE_HASHED))
     for line in virus_defs:
         if FILE_HASHED == line.strip():
             print("[!]Malware Detected[!] | File name: {}".format(f))
             infected_list.append(f)
  • Why do you keep on scanning the rest of the virus definitions if the file is already known to be infected?
    else:
          pass
      except Exception as e:
          print("Could not read file | Error: {}".format(e))
   except:
      pass
  • Why are you swallowing all the exceptions here? What is the potential benefit of this?

    print("Infected files found: {}".format(infected_list))
    deleteornot = str(input("Would you like to delete the infected files (y/n): "))
    
  • input() returns a str , so why convert it?

    if deleteornot.upper() == "Y":
        for infected in infected_list:
            os.remove(infected)
            print("File removed: {}".format(infected))
    else:
        print("Executed with exit code 0")
  • What's the reason for this? You are not actually using an exit code; the output is a lie.
       os.system("PAUSE")
  • Not cross platform.
 Scan()
\$\endgroup\$

Not the answer you're looking for? Browse other questions tagged or ask your own question.