Simplified encrypt.py

encrypt.py

-##############################################################################
-#                                                                            #
-# Script 2 of the Klasureinsicht for all platforms                           #
-# Author: Amrita Deb <Deb@itc.rwth-aachen.de>                                #
-#                                                                            #
-# This scripts creates encrypted copies of the watermarked PDFs created by   #
-# watermark.py as well as a csv files storing passwords for each file        #
-#                                                                            #
-##############################################################################
+#!/usr/bin/env python
+
+"""Creates encrypted copies of PDFs
+
+This scripts creates encrypted copies of the watermarked PDFs created by
+watermark.py as well as a csv files storing passwords for each file
+
+Author: Amrita Deb <Deb@itc.rwth-aachen.de>
+"""
 
 import pikepdf
-import os, math, string, struct,random, csv, pwgen, sys, argparse, time
-
-##############################################################################
-#                                                                            #
-# Encryption function                                                        #
-#                                                                            #
-# Input params: filepath of the watermarked pdfs, filepath of the            #
-#               encrypted pdf, randomly generated 8-character password       #
-# Output: Encrypted file in pdfs_encrypted folder                            #
-#                                                                            #
-##############################################################################
-
-def encrypt(inFile, outFile, password):
-    pdf = pikepdf.Pdf.open(inFile)    
-    pdf.save(outFile, encryption=pikepdf.Encryption(owner=password, user=password, R=4))
+import os
+import csv
+import pwgen
+import argparse
+import sys
+import time
+
+import utils.matnum as utils
+
+
+def encrypt(pdf_file, enc_file, password):
+    """PDF encryption
+
+    Args:
+        pdf_file (str): path to PDF
+        enc_file (str): path of the encrypted pdf
+        password (str): password
+    """
+
+    pdf = pikepdf.Pdf.open(pdf_file)
+    enc = pikepdf.Encryption(owner=password, user=password, R=4)
+    pdf.save(enc_file, encryption=enc)
     pdf.close()
 
-##############################################################################
-#                                                                            #
-# Main function                                                              #
-#                                                                            #
-# 1) Lists all PDFs to be encrypted from ./pdfs_watermarked folder           #
-# 2) Encrypt pdf witn randomly generated 8 character long password           #
-# 3) Prepare a csv file that contains matriculation number and password      #
-#                                                                            #
-##############################################################################
-if __name__ == '__main__':
-    #Defining parameters
-    parser = argparse.ArgumentParser(description='''
-      prepares batch upload to Moodle via assignment module. 
-      PDFs in folder 'in' are moved to folder 'tmp' with a certain folder structure and finally zipped to 'out'. 
-      Attention: zip-archive 'out' will be overwritten in the following!
 
+def main(args):
+    """Main function
+
+    1) Lists all PDFs to be encrypted from input folder
+    2) Encrypt pdf witn randomly generated 8 character long password
+    3) Prepare a csv file that contains matriculation number and password
+    """
+
+    # Argument handling
+    parser = argparse.ArgumentParser(description='''
+      PDFs in input folder are encrypted and stored in output folder.
+      Alongside with a CSV file mapping passwords to each PDF.
       ''')
-    parser.add_argument("-i", "--infolder", default="./pdfs_watermarked", 
+    parser.add_argument("-i", "--infolder", default="./pdfs_watermarked",
         help="Input folder with watermarked PDFs. Default: ./pdfs_watermarked")
 
     parser.add_argument("-o", "--outfolder", default="./pdfs_encrypted",
         help="Output folder of the encrypted PDFs Default: ./pdfs_encrypted")
 
-    parser.add_argument("-p", "--password", default=" ",
-        help="Common password for encrypted PDFs Default: '' will be changed to a 8character randomly generated password")
-    
-    parser.add_argument("-w", "--passwordout", default=" ",
-        help="separate folder for the CSV file containing passwords. This required only for batch.py")
+    parser.add_argument("-p", "--password", default="",
+        help="Common password for all encrypted PDFs. Default: '' will be changed to a 8 character randomly generated password")
 
-    args = parser.parse_args()
+    parser.add_argument("-w", "--passwordout", default="",
+        help="Output path for CSV file. Default: '' will be changed to [outfolder]/passwords.csv.")
+
+    args = parser.parse_args(args)
     infolder = args.infolder
     outfolder = args.outfolder
-    
-    #Empty 'out' folder
-    for root, dirs, files in os.walk(outfolder):
-        for f in files:
-            os.unlink(os.path.join(root, f))
-        for d in dirs:
-            shutil.rmtree(os.path.join(root, d))
-    
-    #List all PDFs
-    pdf_folder = os.listdir(infolder)
-    pdf_files = [_ for _ in pdf_folder if _[-4:] == ".pdf"]
-    print("Available PDFs to be encrypted:\n")
-    for pdffile in pdf_files:
-        print(pdffile)
-    print('\n')
-    pwdfileinputs = []
+
+    if args.passwordout == "":
+        password_file = os.path.join(outfolder, 'passwords.csv')
+    else:
+        password_file = args.passwordout
+
+    # List all PDFs
     starttime = time.time()
-    for i in pdf_files:
-        print('Encrypting '+i+' ...')
-        filename = i.split('.',1)[0]
-        matnum = filename.split('_',1)
-        if args.password==" ":
+    pdf_folder = os.listdir(infolder)
+    pdf_files = [_ for _ in pdf_folder
+                 if _.endswith(".pdf") and utils.check_matnum(_[0:6])]
+    print("""
+Available PDFs to be encrypted:
+- {}
+
+Files in output folder {} will be overwritten during this process.
+    """.format("\n- ".join(pdf_files), outfolder))
+
+    # Encrypt all PDFs in input folder
+    csv_lines = []
+    enc_files = []
+    for pdf_file in pdf_files:
+        # PDF file has to start with 6 digit mat number
+        matnum = utils.get_matnum(pdf_file)
+
+        # Generate random password if common password not given
+        if args.password == "":
             password = pwgen.pwgen(8)
         else:
             password = args.password
-        encrypt(infolder+'/'+i, outfolder+'/'+filename+'_aes.pdf', password)
-        pwdfileinputs.append(matnum[0]+', '+password)
-        print('Encryption completed for '+i+'\n')
-    endtime = time.time()    
-    print('Recording password..')
-    if args.passwordout == " ":
-        with open(outfolder+'/passwords.csv', mode='w') as password_file:
-            pwdfile_writer = csv.writer(password_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
-            for eachinput in pwdfileinputs:
-                pwdfile_writer.writerow([eachinput.split(',',1)[0],eachinput.split(',',1)[1]])
-    else:
-        with open(args.passwordout+'/passwords.csv', mode='w') as password_file:
-            pwdfile_writer = csv.writer(password_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
-            for eachinput in pwdfileinputs:
-                pwdfile_writer.writerow([eachinput.split(',',1)[0],eachinput.split(',',1)[1]])
-        
-    
-    print('\nThe encrypted files are available in the folder'+ outfolder+'. \nThe passwords for the encrypted files are stored in passwords.csv in the '+args.passwordout+' folder\n')
-    print(f'\nTime taken: {endtime-starttime:.2f}s\n') 
+
+        in_file = os.path.join(infolder, pdf_file)
+        enc_file = os.path.splitext(pdf_file)[0] + '_aes.pdf'
+        enc_file = os.path.join(outfolder, enc_file)
+        encrypt(in_file, enc_file, password)
+
+        # Save matnum password mapping to be stored in CSV later
+        csv_lines.append([matnum, password])
+        enc_files.append(enc_file)
+
+    # Store matnum password mappings in CSV file
+    print('Saving passwords')
+    with open(password_file, mode='w') as password_handle:
+        writer = csv.writer(password_handle, delimiter=',', quotechar='"',
+                            quoting=csv.QUOTE_MINIMAL)
+
+        for csv_line in csv_lines:
+            writer.writerow(csv_line)
+
+    # Print status
+    endtime = time.time()
+    print("""All PDFs are encrypted and can be found in {} folder.
+CSV file is stored in {}.
+
+Time taken: {:.2f}s
+    """.format(outfolder, password_file, endtime-starttime))
+
+    return enc_files, csv_lines
+
+
+if __name__ == '__main__':
+    main(sys.argv[1:])

tests/test_encrypt.py

+import unittest
+import time
+import os
+import tempfile
+import shutil
+
+
+class MainTest(unittest.TestCase):
+    def setUp(self):
+        self.tic = time.time()  # todo this is sooo ugly
+        self.test_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        self.toc = time.time()
+        t = self.toc - self.tic
+        print('Time: %.3f' % (t))
+
+        # Clean up
+        shutil.rmtree(self.test_dir)
+
+    def test_encrypt_pdfs(self):
+        import encrypt
+
+        expected_files = ['123456_Nachname_aes.pdf', '456789_Lastname_aes.pdf',
+                          'passwords.csv']
+
+        # Prepare parameter
+        in_dir = './pdfs'
+
+        out_dir = os.path.join(self.test_dir, 'out')
+        os.mkdir(out_dir)
+
+        # Encrypt files
+        encrypt.main(["-i", in_dir, "-o", out_dir])
+
+        created_files = os.listdir(out_dir)
+        created_files.sort()
+        self.assertEqual(expected_files, created_files)