Commit 57b52b42 authored by Christian Rohlfing's avatar Christian Rohlfing
Browse files

Simplified encrypt.py

parent 7c86314c
##############################################################################
# #
# 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",
help="Input folder with watermarked PDFs. Default: ./pdfs_watermarked")
......@@ -49,56 +53,74 @@ if __name__ == '__main__':
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("-p", "--password", default="",
help="Common password for all encrypted PDFs. Default: '' will be changed to a 8 character 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("-w", "--passwordout", default="",
help="Output path for CSV file. Default: '' will be changed to [outfolder]/passwords.csv.")
args = parser.parse_args()
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))
if args.passwordout == "":
password_file = os.path.join(outfolder, 'passwords.csv')
else:
password_file = args.passwordout
#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 = []
# 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')
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('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("""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))
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')
return enc_files, csv_lines
if __name__ == '__main__':
main(sys.argv[1:])
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)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment