<commit_before>import types
import logging
from banal import ensure_list
from normality import stringify
from balkhash.utils import safe_fragment
from email.utils import parsedate_to_datetime, getaddresses
from normality import safe_filename, ascii_text
from followthemoney.types import registry

from ingestors.support.html import HTMLSupport
from ingestors.support.cache import CacheSupport
from ingestors.support.temp import TempFileSupport

log = logging.getLogger(__name__)


class EmailIdentity(object):

    def __init__(self, manager, name, email):
        self.email = ascii_text(stringify(email))
        self.name = stringify(name)
        if not registry.email.validate(self.email):
            self.email = None
        if registry.email.validate(self.name):
            self.email = self.email or ascii_text(self.name)
            self.name = None

        # This should be using formataddr, but I cannot figure out how
        # to use that without encoding the name.
        self.label = None
        if self.name is not None and self.email is not None:
            self.label = '%s <%s>' % (self.name, self.email)
        elif self.name is None and self.email is not None:
            self.label = self.email
        elif self.email is None and self.name is not None:
            self.label = self.name

        self.entity = None
        if self.email is not None:
            key = self.email.lower().strip()
            fragment = safe_fragment(self.label)
            self.entity = manager.make_entity('LegalEntity')
            self.entity.make_id(key)
            self.entity.add('name', self.name)
            self.entity.add('email', self.email)
            manager.emit_entity(self.entity, fragment=fragment)


class EmailSupport(TempFileSupport, HTMLSupport, CacheSupport):
    """Extract metadata from email messages."""

    def ingest_attachment(self, entity, name, mime_type, body):
        has_body = body is not None and len(body)
        if stringify(name) is None and not has_body:
            # Hello, Outlook.
            return

        file_name = safe_filename(name, default='attachment')
        file_path = self.make_work_file(file_name)
        with open(file_path, 'wb') as fh:
            if isinstance(body, str):
                body = body.encode('utf-8')
            if body is not None:
                fh.write(body)

        checksum = self.manager.store(file_path, mime_type=mime_type)
        file_path.unlink()

        child = self.manager.make_entity('Document', parent=entity)
        child.make_id(name, checksum)
        child.add('contentHash', checksum)
        child.add('fileName', name)
        child.add('mimeType', mime_type)
        self.manager.queue_entity(child)

    def get_header(self, msg, *headers):
        values = []
        for header in headers:
            try:
                for value in ensure_list(msg.get_all(header)):
                    values.append(value)
            except TypeError as te:
                log.warning("Failed to parse [%s]: %s", header, te)
        return values

    def get_dates(self, msg, *headers):
        dates = []
        for value in self.get_header(msg, *headers):
            try:
                dates.append(parsedate_to_datetime(value))
            except Exception:
                log.warning("Failed to parse: %s", value)
        return dates

    def get_identities(self, values):
        values = [v for v in ensure_list(values) if v is not None]
        for (name, email) in getaddresses(values):
            yield EmailIdentity(self.manager, name, email)

    def get_header_identities(self, msg, *headers):
        yield from self.get_identities(self.get_header(msg, *headers))

    def apply_identities(self, entity, identities, eprop=None, lprop=None):
        if isinstance(identities, types.GeneratorType):
            identities = list(identities)
        for identity in ensure_list(identities):
            if eprop is not None:
                entity.add(eprop, identity.entity)
            if lprop is not None:
                entity.add(lprop, identity.label)
            entity.add('namesMentioned', identity.name)
            entity.add('emailMentioned', identity.email)

    def _clean_message_id(self, message_id):
        message_id = stringify(message_id)
        if message_id is not None:
            message_id = message_id.strip()
            if len(message_id) > 4:
                return message_id

    def resolve_message_ids(self, entity):
        ctx = self.manager.queue.dataset
        for message_id in entity.get('messageId'):
            message_id = self._clean_message_id(message_id)
            if message_id is None:
                continue
            key = self.cache_key('msid', ctx, message_id)
            self.set_cache_value(key, entity.id)
            rev_key = self.cache_key('mrid', ctx, message_id)
            for entity_id in self.get_cache_set(rev_key):
                email = self.manager.make_entity('Email')
                email.id = entity_id
                email.add('inReplyToEmail', message_id)
                fragment = safe_fragment(message_id)<commit_msg>Fix email thread response ID<commit_after>import types
import logging
from banal import ensure_list
from normality import stringify
from balkhash.utils import safe_fragment
from email.utils import parsedate_to_datetime, getaddresses
from normality import safe_filename, ascii_text
from followthemoney.types import registry

from ingestors.support.html import HTMLSupport
from ingestors.support.cache import CacheSupport
from ingestors.support.temp import TempFileSupport

log = logging.getLogger(__name__)


class EmailIdentity(object):

    def __init__(self, manager, name, email):
        self.email = ascii_text(stringify(email))
        self.name = stringify(name)
        if not registry.email.validate(self.email):
            self.email = None
        if registry.email.validate(self.name):
            self.email = self.email or ascii_text(self.name)
            self.name = None

        # This should be using formataddr, but I cannot figure out how
        # to use that without encoding the name.
        self.label = None
        if self.name is not None and self.email is not None:
            self.label = '%s <%s>' % (self.name, self.email)
        elif self.name is None and self.email is not None:
            self.label = self.email
        elif self.email is None and self.name is not None:
            self.label = self.name

        self.entity = None
        if self.email is not None:
            key = self.email.lower().strip()
            fragment = safe_fragment(self.label)
            self.entity = manager.make_entity('LegalEntity')
            self.entity.make_id(key)
            self.entity.add('name', self.name)
            self.entity.add('email', self.email)
            manager.emit_entity(self.entity, fragment=fragment)


class EmailSupport(TempFileSupport, HTMLSupport, CacheSupport):
    """Extract metadata from email messages."""

    def ingest_attachment(self, entity, name, mime_type, body):
        has_body = body is not None and len(body)
        if stringify(name) is None and not has_body:
            # Hello, Outlook.
            return

        file_name = safe_filename(name, default='attachment')
        file_path = self.make_work_file(file_name)
        with open(file_path, 'wb') as fh:
            if isinstance(body, str):
                body = body.encode('utf-8')
            if body is not None:
                fh.write(body)

        checksum = self.manager.store(file_path, mime_type=mime_type)
        file_path.unlink()

        child = self.manager.make_entity('Document', parent=entity)
        child.make_id(name, checksum)
        child.add('contentHash', checksum)
        child.add('fileName', name)
        child.add('mimeType', mime_type)
        self.manager.queue_entity(child)

    def get_header(self, msg, *headers):
        values = []
        for header in headers:
            try:
                for value in ensure_list(msg.get_all(header)):
                    values.append(value)
            except TypeError as te:
                log.warning("Failed to parse [%s]: %s", header, te)
        return values

    def get_dates(self, msg, *headers):
        dates = []
        for value in self.get_header(msg, *headers):
            try:
                dates.append(parsedate_to_datetime(value))
            except Exception:
                log.warning("Failed to parse: %s", value)
        return dates

    def get_identities(self, values):
        values = [v for v in ensure_list(values) if v is not None]
        for (name, email) in getaddresses(values):
            yield EmailIdentity(self.manager, name, email)

    def get_header_identities(self, msg, *headers):
        yield from self.get_identities(self.get_header(msg, *headers))

    def apply_identities(self, entity, identities, eprop=None, lprop=None):
        if isinstance(identities, types.GeneratorType):
            identities = list(identities)
        for identity in ensure_list(identities):
            if eprop is not None:
                entity.add(eprop, identity.entity)
            if lprop is not None:
                entity.add(lprop, identity.label)
            entity.add('namesMentioned', identity.name)
            entity.add('emailMentioned', identity.email)

    def _clean_message_id(self, message_id):
        message_id = stringify(message_id)
        if message_id is not None:
            message_id = message_id.strip()
            if len(message_id) > 4:
                return message_id

    def resolve_message_ids(self, entity):
        ctx = self.manager.queue.dataset
        for message_id in entity.get('messageId'):
            message_id = self._clean_message_id(message_id)
            if message_id is None:
                continue
            key = self.cache_key('msid', ctx, message_id)
            self.set_cache_value(key, entity.id)
            rev_key = self.cache_key('mrid', ctx, message_id)
            for response_id in self.get_cache_set(rev_key):
                email = self.manager.make_entity('Email')
                email.id = response_id
                email.add('inReplyToEmail', entity.id)
                fragment = safe_fragment(message_id)
{'fraction_non_alphanumeric': 0.09093550261475886, 'fraction_numerical': 0.0003873716831299632, 'mean_word_length': 2.2804955527318933, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, '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': '3277710', 'n_tokens_mistral': 2921, 'n_tokens_neox': 2841, 'n_words': 769}