DOCX Helper Functions

docx_helpers.py

""" Helpers for working with docx files.

Currently includes a couple functions for adding additional permissions to
protect contents of the document.
"""
import string
import random

import lxml.etree
from docx.oxml.settings import CT_Settings


def generate_dummy_string(lenght: int) -> str:
    """ Generates a random string of the specified length. """
    return ''.join([
        random.choice(
            string.ascii_letters
            + string.digits
            + string.punctuation
        )
        for _ in range(lenght)
    ])


def add_namespace(namespace: str, name: str) -> str:
    """Adds the namespace to the XML name."""
    return '{%s}%s' % (namespace, name)


def add_ro_recommendation(doc) -> None:
    """Adds an XML element telling docx to recommend read-only mode when
    opening the document.
    """

    root: CT_Settings = doc.settings.element
    nsmap: dict[str, str] = root.nsmap
    prefix: str = nsmap['w']

    _ = lxml.etree.SubElement(
        root,
        add_namespace(prefix, 'writeProtection'),
        attrib={
            add_namespace(prefix, 'recommended'): '1',
        },
    )


def make_doc_read_only(doc) -> None:
    """Makes the document read-only."""

    root: CT_Settings = doc.settings.element
    nsmap: dict[str, str] = root.nsmap
    prefix: str = nsmap['w']

    _ = lxml.etree.SubElement(
        root,
        add_namespace(prefix, 'documentProtection'),
        attrib={
            add_namespace(prefix, 'recommended'): '1',
            add_namespace(prefix, 'edit'): "readOnly",
            add_namespace(prefix, 'enforcement'): "1",
            add_namespace(prefix, 'cryptProviderType'): "rsaAES",
            add_namespace(prefix, 'cryptAlgorithmClass'): "hash",
            add_namespace(prefix, 'cryptAlgorithmType'): "typeAny",  # noqa: E501
            add_namespace(prefix, 'cryptAlgorithmSid'): "14",
            add_namespace(prefix, 'cryptSpinCount'): "100000",
            add_namespace(
                prefix, 'hash'
            ): generate_dummy_string(86) + '==',  # Dummy hash
            add_namespace(prefix, 'salt'):
            generate_dummy_string(22) + '==',  # Dummy salt
        },
    )