OrdersGuru_signer/helpers/xmlhandler.py

from lxml import etree
from helpers.xadessigner import XAdESSigner, XaDESSignerException
from copy import deepcopy
import unicodedata
import string


class XMLHandler:
    def __init__(self, xml_file, xml_reference_type):
        self.xml_file = xml_file
        self.xml_tree = self.__parse_xml(xml_file)
        self.xml_reference_type = xml_reference_type

    def __parse_xml(self, xml_file):
        try:
            xml_tree = etree.parse(xml_file)
            return xml_tree
        except:
            raise XMLHandlerException('Input XML is invalid.')

    def valid(self, xml_schema):
        try:
            xml_schema_tree = etree.XMLSchema(etree.parse(xml_schema))
            return xml_schema_tree.validate(self.xml_tree)
        except:
            return False

    def sign(self, certificate_file, certificate_password):
        try:
            xml_signer = XAdESSigner(deepcopy(self.xml_tree.getroot()), self.xml_reference_type, certificate_file)
            xml_signed_tree = xml_signer.sign(certificate_password)
            with open(self.xml_file, 'wb') as signed_xml:
                etree.ElementTree(xml_signed_tree).write(signed_xml, encoding='utf-8', xml_declaration=True)
        except IOError:
            raise XMLHandlerException('Problem with writing signed XML file.')
        except XaDESSignerException as e:
            raise XMLHandlerException(str(e))
        except:
            raise XMLHandlerException('Something went wrong with writing signed XML file.')

    def get_general_filename(self):
        default_filename = 'document'
        xml_root = self.xml_tree.getroot()
        if xml_root is not None and len(xml_root) > 0 and len(xml_root[0]) > 0 and len(xml_root[0][0]) > 1:
            return self.__normalize_filename(xml_root[0][0][1].text.strip(), default_filename)
        return default_filename

    def __normalize_filename(self, filename, default_filename='document'):
        if len(filename) == 0:
            return default_filename
        valid_chars = '-_%s%s' % (string.ascii_letters, string.digits)
        replacement_char = '_'
        clean_name = unicodedata.normalize('NFKD', filename).encode('ascii', 'ignore').decode('utf-8')
        return ''.join(c.lower() if c in valid_chars else replacement_char for c in clean_name)


class XMLHandlerException(Exception):
    pass