Encrypting your files using a USB stick key (Python)

Cyber security is an increasingly important aspect to our digital lives, and people really underestimate how easily their information can be stolen. In light of this, why not set up a cool little system of your own in python that can encrypt your files! In this post I will go through how you can create a secure location on you windows pc, which will encrypt your important files, unless you plug in one of your USB sticks.

Now I won’t pretend that this is the most secure thing you can do to protect yourself, but every little helps and it’s very simple. All you will need is a USB stick, preferably with nothing on it, but that’s not essential.

To start with you will want to plug in your USB stick, and make sure it works fine, and clear off any useless or unwanted files. When it appears on your computer, in the side menu, right click on it and click on properties. In here you can give it a name, and in my case I will just call it ‘KEY’. The other thing you want to do before we jump into the code is to choose where on your computer you want to keep these secure files. I have just made a folder on my desktop called SecureData, quick and easy.

So into the code! We need three libraries for this project, two of which will require a pip install.

pip install wmi
pip install cryptography
import os
import wmi
from cryptography.fernet import Fernet

Then you want to define a variable that holds the string name that you gave to your USB stick, and another variable that holds the system path to your secure files folder.

my_key = 'KEY'
files = r'YOUR PATH HERE'

Now we instantiate an object of ‘wmi’ which we imported. This is the library that allows us to scan the ports of the computer.

c = wmi.WMI()

That’s all the setup we need for now. The next part is broken down into only four different functions. The first of which is to check if our ‘key’ is plugged in or not.

def check_for_key():
   for disk in c.Win32_LogicalDisk():
      if disk.VolumeName==my_key:
         return disk

Very simple and easy. This just iterates over any disk drives, including your internal disks, but only returns a value if the drive name matches they name you gave it.

The first part of the encryption process is that we need to generate a unique encryption key. This is how the data we know an understand is converted into utter gibberish. So in this function we want to read from our USB stick, see if there is a key already, or if not, create a new one. To do this is very simple, we will use a ‘try’ block to try and load the key, and in the case there isn’t one, in the ‘except’ block we will create a new one.

def load_key(usbDisk):
   port = usbDisk.DeviceID
   try:
      print('Trying to find key...')
      with open(f'{port}\\encryptionKey.key','rb') as encryptKey:
         key = encryptKey.read()
         print('Key Found')
   except:
      print('Key not found... Creating a new key')
      key = Fernet.generate_key()
      with open(f'{port}\\encryptionKey.key','wb') as encryptKey:
         encryptKey.write(key)
   return key

So now we have our USB port scanner, and an encryption key ready to secure our data. The next two functions we need to make are the encryption and decryption functions. These could definitely be condensed into one function, but for simplicity its easier to keep them separate.

Starting with the encryption function:

def encryptFiles(key,directory):
   files = os.listdir(directory)
   cipher = Fernet(key)
   global state
   state = 'encrypted'
   for file in files:
      with open(f'{directory}\{file}','rb') as old:
         original = old.read()
      encrypted = cipher.encrypt(original)
      with open(f'{directory}\{file}','wb') as old:
         old.write(encrypted)

So just for clarity, the parameters we are passing in are our key, and then the directory is the path to our secure files folder we defined at the start. The global declaration of the variable state will become more clear later. Then we open each file in our directory, and read it in. We encrypt the contents, then open the file again in the ‘write’ mode to replace the contents with our encrypted data.

So now the decryption function, which is almost identical:

def decryptFiles(key,directory):
   files = os.listdir(directory)
   cipher = Fernet(key)
   global state
   state = 'decrypted'
   for file in files:
      with open(f'{directory}\{file}','rb') as old:
         encrypted = old.read()
      decrypted = cipher.decrypt(encrypted)
      with open(f'{directory}\{file}','wb') as old:
         old.write(decrypted)

And that is the only functions we need to write! We just have one more small part to add to the bottom of our code, which is going to put these functions to use.

state = 'decrypted'
if __name__=='__main__':
   while True:
      disk = check_for_key()
      try:
         key = load_key(disk)
      except:
         print('No Key Available')
      if disk!=None:
         current_state = 'decrypted'
         if current_state!=state:
            decryptFiles(key,files)
         else:
            current_state = 'encrypted'
            if current_state!=state:
               encryptFiles(key,files)

So in this final section, we are just making sure that the files only get encrypted or decrypted when the ‘key’ is plugged in or removed, instead of constantly, which would put a real strain on the computer, and likely make the files impossible to open.

It is important to note that the way this last part is coded means that when you start running the script the files need to be in their normal format (decrypted) and that your key is plugged in. After that it will work exactly as you want. Another note is that not all file types respond well to being encrypted, it works best with plain text formats and csv file types. PDF’s in particular may get corrupted, so be careful which files you use this with, it’s on you!

I hope you enjoy experimenting with encrypting your files, and that this is a helpful introduction to using python for keeping your important data secure!