<commit_before>import shutil
import logging
import zipfile
import pathlib
from lxml import etree
from pprint import pprint  # noqa
from normality import safe_filename
from followthemoney import model

from ingestors.ingestor import Ingestor
from ingestors.support.temp import TempFileSupport
from ingestors.support.timestamp import TimestampSupport
from ingestors.support.email import EmailSupport, EmailIdentity
from ingestors.exc import ProcessingException
from ingestors.util import safe_string

log = logging.getLogger(__name__)
MIME = 'application/xml+opfmessage'


class OPFParser(object):

    def parse_xml(self, file_path):
        parser = etree.XMLParser(huge_tree=True)
        try:
            return etree.parse(file_path.as_posix(), parser)
        except etree.XMLSyntaxError:
            # probably corrupt
            raise TypeError()


class OutlookOLMArchiveIngestor(Ingestor, TempFileSupport, OPFParser):
    MIME_TYPES = []
    EXTENSIONS = ['olm']
    SCORE = 10
    EXCLUDE = ['com.microsoft.__Messages']

    def extract_file(self, zipf, name):
        path = pathlib.Path(name)
        base_name = safe_filename(path.name)
        out_file = self.make_work_file(base_name)
        with open(out_file, 'w+b') as outfh:
            try:
                with zipf.open(name) as infh:
                    shutil.copyfileobj(infh, outfh)
            except KeyError:
                log.warning("Cannot load zip member: %s", name)
        return out_file

    def extract_hierarchy(self, entity, name):
        foreign_id = pathlib.PurePath(entity.id)
        path = pathlib.Path(name)
        for name in path.as_posix().split('/'):
            foreign_id = foreign_id.joinpath(name)
            if name in self.EXCLUDE:
                continue
            entity = self.manager.make_entity('Folder', parent=entity)
            entity.make_id(foreign_id.as_posix())
            self.manager.emit_entity(entity)
        return entity

    def extract_attachment(self, zipf, message, attachment):
        url = attachment.get('OPFAttachmentURL')
        name = attachment.get('OPFAttachmentName')
        name = name or attachment.get('OPFAttachmentContentID')
        child = self.manager.make_entity('Document', parent=message)
        if url is not None:
            file_path = self.extract_file(zipf, url)
            checksum = self.manager.archive_store(file_path)
            child.make_id(name, checksum)
            child.add('fileName', attachment.get('OPFAttachmentName'))<commit_msg>Set folder names while ingesting OLM mailboxes<commit_after>import shutil
import logging
import zipfile
import pathlib
from lxml import etree
from pprint import pprint  # noqa
from normality import safe_filename
from followthemoney import model

from ingestors.ingestor import Ingestor
from ingestors.support.temp import TempFileSupport
from ingestors.support.timestamp import TimestampSupport
from ingestors.support.email import EmailSupport, EmailIdentity
from ingestors.exc import ProcessingException
from ingestors.util import safe_string

log = logging.getLogger(__name__)
MIME = 'application/xml+opfmessage'


class OPFParser(object):

    def parse_xml(self, file_path):
        parser = etree.XMLParser(huge_tree=True)
        try:
            return etree.parse(file_path.as_posix(), parser)
        except etree.XMLSyntaxError:
            # probably corrupt
            raise TypeError()


class OutlookOLMArchiveIngestor(Ingestor, TempFileSupport, OPFParser):
    MIME_TYPES = []
    EXTENSIONS = ['olm']
    SCORE = 10
    EXCLUDE = ['com.microsoft.__Messages']

    def extract_file(self, zipf, name):
        path = pathlib.Path(name)
        base_name = safe_filename(path.name)
        out_file = self.make_work_file(base_name)
        with open(out_file, 'w+b') as outfh:
            try:
                with zipf.open(name) as infh:
                    shutil.copyfileobj(infh, outfh)
            except KeyError:
                log.warning("Cannot load zip member: %s", name)
        return out_file

    def extract_hierarchy(self, entity, name):
        foreign_id = pathlib.PurePath(entity.id)
        path = pathlib.Path(name)
        for name in path.as_posix().split('/'):
            foreign_id = foreign_id.joinpath(name)
            if name in self.EXCLUDE:
                continue
            entity = self.manager.make_entity('Folder', parent=entity)
            entity.add('fileName', name)
            entity.make_id(foreign_id.as_posix())
            self.manager.emit_entity(entity)
        return entity

    def extract_attachment(self, zipf, message, attachment):
        url = attachment.get('OPFAttachmentURL')
        name = attachment.get('OPFAttachmentName')
        name = name or attachment.get('OPFAttachmentContentID')
        child = self.manager.make_entity('Document', parent=message)
        if url is not None:
            file_path = self.extract_file(zipf, url)
            checksum = self.manager.archive_store(file_path)
            child.make_id(name, checksum)
            child.add('fileName', attachment.get('OPFAttachmentName'))
{'fraction_non_alphanumeric': 0.08688845401174168, 'fraction_numerical': 0.0007827788649706458, 'mean_word_length': 2.8370870870870872, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 6}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '814082', 'n_tokens_mistral': 1505, 'n_tokens_neox': 1468, 'n_words': 347}