def decrypt_pdf(update, context):
result = check_back_user_data(update, context)
if result is not None:
return result
_ = set_lang(update, context)
message = update.effective_message
message.reply_text(
_("Decrypting your PDF file"), reply_markup=ReplyKeyboardRemove()
)
with tempfile.NamedTemporaryFile() as tf:
# Download file
user_data = context.user_data
file_id, file_name = user_data[PDF_INFO]
pdf_file = context.bot.get_file(file_id)
pdf_file.download(custom_path=tf.name)
pdf_reader = None
try:
pdf_reader = PdfFileReader(open(tf.name, "rb"))
except PdfReadError:
message.reply_text(
_("Your PDF file seems to be invalid and I couldn't open and read it")
)
if pdf_reader is not None:
if not pdf_reader.isEncrypted:
message.reply_text(_("Your PDF file is not encrypted"))
else:
try:
if pdf_reader.decrypt(message.text) == 0:
message.reply_text(
_(
"The decryption password is incorrect, try to send it again"
)
)
return WAIT_DECRYPT_PW
pdf_writer = PdfFileWriter()
for page in pdf_reader.pages:
pdf_writer.addPage(page)
write_send_pdf(update, context, pdf_writer, file_name, "decrypted")
except NotImplementedError:
message.reply_text(
_(
"Your PDF file is encrypted with a method that I cannot decrypt"
)
)
# Clean up memory
if user_data[PDF_INFO] == file_id:
del user_data[PDF_INFO]
return ConversationHandler.END
def merge_pdf(update, context):
_ = set_lang(update, context)
update.effective_message.reply_text(_("Merging your PDF files"),
reply_markup=ReplyKeyboardRemove())
# Setup temporary files
user_data = context.user_data
file_ids = user_data[MERGE_IDS]
file_names = user_data[MERGE_NAMES]
temp_files = [tempfile.NamedTemporaryFile() for _ in range(len(file_ids))]
merger = PdfFileMerger()
# Merge PDF files
for i, file_id in enumerate(file_ids):
file_name = temp_files[i].name
file = context.bot.get_file(file_id)
file.download(custom_path=file_name)
try:
merger.append(open(file_name, "rb"))
except PdfReadError:
update.effective_message.reply_text(
_("I can't merge your PDF files as I couldn't open and read \"{}\". "
"Ensure that it is not encrypted").format(file_names[i]))
return ConversationHandler.END
# Send result file
write_send_pdf(update, context, merger, "files.pdf", "merged")
# Clean up memory and files
if user_data[MERGE_IDS] == file_ids:
del user_data[MERGE_IDS]
if user_data[MERGE_NAMES] == file_names:
del user_data[MERGE_NAMES]
for tf in temp_files:
tf.close()
return ConversationHandler.END
def add_wmk(update, context):
if not check_user_data(update, context, WMK_ID):
return ConversationHandler.END
_ = set_lang(update, context)
update.effective_message.reply_text(
_("Adding the watermark onto your PDF file"),
reply_markup=ReplyKeyboardRemove())
# Setup temporary files
temp_files = [tempfile.NamedTemporaryFile() for _ in range(2)]
src_fn, wmk_fn = [x.name for x in temp_files]
user_data = context.user_data
src_file_id = user_data[WMK_ID]
wmk_file_id = update.effective_message.document.file_id
src_reader = open_pdf(update, context, src_file_id, src_fn)
if src_reader is not None:
wmk_reader = open_pdf(update, context, wmk_file_id, wmk_fn,
_("watermark"))
if wmk_reader is not None:
# Add watermark
pdf_writer = PdfFileWriter()
for page in src_reader.pages:
page.mergePage(wmk_reader.getPage(0))
pdf_writer.addPage(page)
# Send result file
write_send_pdf(update, context, pdf_writer, "file.pdf",
"watermarked")
# Clean up memory and files
if user_data[WMK_ID] == src_file_id:
del user_data[WMK_ID]
for tf in temp_files:
tf.close()
return ConversationHandler.END
def split_pdf(update: Update, context: CallbackContext) -> int:
result = check_back_user_data(update, context)
if result is not None:
return result
_ = set_lang(update, context)
message = update.effective_message
split_range = message.text
if not PageRange.valid(split_range):
message.reply_text(
_(
"The range is invalid. Try again",
reply_markup=get_back_markup(update, context),
))
return WAIT_SPLIT_RANGE
message.reply_text(_("Splitting your PDF file"),
reply_markup=ReplyKeyboardRemove())
with tempfile.NamedTemporaryFile() as tf:
user_data = context.user_data
file_id, file_name = user_data[PDF_INFO]
pdf_reader = open_pdf(update, context, file_id, tf.name)
if pdf_reader is not None:
merger = PdfFileMerger()
merger.append(pdf_reader, pages=PageRange(split_range))
write_send_pdf(update, context, merger, file_name, "split")
# Clean up memory
if user_data[PDF_INFO] == file_id:
del user_data[PDF_INFO]
return ConversationHandler.END
