"""Extract and optionallly defang Indicators of Compromise (IOCs) from text.

All methods return iterator objects, not lists. If for some reason you need
a list, do e.g.: ``list(extract_iocs(my_data))``.

Otherwise, you can iterate over the objects (e.g. in a ``for`` loop) normally.
Each object yielded from the generators will by of type :class:`str`.
"""
import re
import sys
import itertools
import argparse
try:
    # python3
    from urllib.parse import urlparse
    unicode = str
except ImportError:
    from urlparse import urlparse

import ipaddress

# Get basic url format, including a few obfuscation techniques, main anchor is the uri scheme
GENERIC_URL_RE = re.compile(r"([fhstu]\w\w?[px]s?(?::\/\/|__?)\x20?\S+?(?:\x20[\/\.][^\.\/\s]\S*?)*)[\.\?>\"!]*(?=\s|$)")

# Get some obfuscated urls, main anchor is brackets around the period
BRACKET_URL_RE = re.compile(r"\b(\S+(?:\x20?[\(\[]\x20?\.\x20?[\]\)]\x20?\S*?)+)[\.\?>\"!]*(?=\s|$)")

# Get some valid obfuscated ip addresses
IPV4_RE = re.compile(r"(?:^|(?![^\d\.]))(?:(?:[1-9]?\d|1\d\d|2[0-4]\d|25[0-5])[\[\(]?\.[\]\)]?){3}(?:[1-9]?\d|1\d\d|2[0-4]\d|25[0-5])(?:(?=[^\d\.])|$)")

# Experimental IPV6 regex, will not catch everything but should be sufficent for now
IPV6_RE = re.compile(r"\b(?:[a-f0-9]{1,4}:|:){2,7}(?:[a-f0-9]{1,4}|:)\b")

EMAIL_RE = re.compile(r"([a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+)")
MD5_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{32})(?:[^a-fA-F\d]|\b)")
SHA1_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{40})(?:[^a-fA-F\d]|\b)")
SHA256_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{64})(?:[^a-fA-F\d]|\b)")
SHA512_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{128})(?:[^a-fA-F\d]|\b)")

# YARA regex
YARA_SPLIT_STR = r"\n[\t\s]*\}[\s\t]*(rule[\t\s][^\r\n]+(?:\{|[\r\n][\r\n\s\t]*\{))"
YARA_PARSE_RE = re.compile(r"^[\t\s]*(rule[\t\s][^\r\n]+(?:\{|[\r\n][\r\n\s\t]*\{).*?condition:.*?\r?\n?[\t\s]*\})[\s\t]*(?:$|\r?\n)", re.MULTILINE | re.DOTALL)


def extract_iocs(data, refang=False):
    """Extract all IOCs.

    Results are returned as an itertools.chain iterable object which
    lazily provides the results of the other extract_* generators.

    :param data: Input text
    :param bool refang: Refang output?
    :rtype: :py:func:`itertools.chain`
    """
    return itertools.chain(
        extract_urls(data, refang=refang),
        extract_ips(data, refang=refang),
        extract_emails(data),
        extract_hashes(data),
        extract_yara_rules(data)
    )

def extract_urls(data, refang=False):
    """Extract URLs.

    :param data: Input text
    :param bool refang: Refang output?
    :rtype: Iterator[:class:`str`]
    """
    for url in GENERIC_URL_RE.finditer(data):
        if refang:
            yield refang_url(url.group(1))
        else:
            yield url.group(1)
    for url in BRACKET_URL_RE.finditer(data):
        if refang:
            yield refang_url(url.group(1))
        else:
            yield url.group(1)

def extract_ips(data, refang=False):
    """Extract IP addresses.

    Includes both IPv4 and IPv6 addresses.

    :param data: Input text
    :param bool refang: Refang output?
    :rtype: :py:func:`itertools.chain`
    """
    return itertools.chain(
        extract_ipv4s(data, refang=refang),
        extract_ipv6s(data),
    )

def extract_ipv4s(data, refang=False):
    """Extract IPv4 addresses.

    :param data: Input text
    :param bool refang: Refang output?
    :rtype: Iterator[:class:`str`]
    """
    for ip_address in IPV4_RE.finditer(data):
        if refang:
            yield refang_ipv4(ip_address.group(0))
        else:
            yield ip_address.group(0)

def extract_ipv6s(data):
    """Extract IPv6 addresses.

    Not guaranteed to catch all valid IPv6 addresses.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for ip_address in IPV6_RE.finditer(data):
        yield ip_address.group(0)

def extract_emails(data):
    """Extract email addresses

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for email in EMAIL_RE.finditer(data):
        yield email.group(0)

def extract_hashes(data):
    """Extract MD5/SHA hashes.

    Results are returned as an itertools.chain iterable object which
    lazily provides the results of the other extract_*_hashes generators.

    :param data: Input text
    :rtype: :py:func:`itertools.chain`
    """
    return itertools.chain(
        extract_md5_hashes(data),
        extract_sha1_hashes(data),
        extract_sha256_hashes(data),
        extract_sha512_hashes(data)
    )

def extract_md5_hashes(data):
    """Extract MD5 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for md5 in MD5_RE.finditer(data):
        yield md5.group(1)

def extract_sha1_hashes(data):
    """Extract SHA1 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for sha1 in SHA1_RE.finditer(data):
        yield sha1.group(1)

def extract_sha256_hashes(data):
    """Extract SHA256 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for sha256 in SHA256_RE.finditer(data):
        yield sha256.group(1)

def extract_sha512_hashes(data):
    """Extract SHA512 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for sha512 in SHA512_RE.finditer(data):
        yield sha512.group(1)

def extract_yara_rules(data):
    """Extract YARA rules.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    yara_rules = re.sub(YARA_SPLIT_STR, "}\r\n\\1", data,
                        re.MULTILINE | re.DOTALL)
    for yara_rule in YARA_PARSE_RE.finditer(yara_rules):
        yield yara_rule.group(1)

def _is_ipv6_url(url):
    """URL network location is an IPv6 address, not a domain.

    :param url: String URL
    :rtype: bool
    """
    # fix urlparse exception
    parsed = urlparse(url)

    # Handle RFC 2732 IPv6 URLs with and without port, as well as non-RFC IPv6 URLs
    if ']:' in parsed.netloc:
        ipv6 = ':'.join(parsed.netloc.split(':')[:-1])
    else:
        ipv6 = parsed.netloc

    try:
        ipaddress.IPv6Address(unicode(ipv6.replace('[', '').replace(']', '')))
    except ValueError:
        return False

    return True

def _refang_common(ioc):
    """Remove artifacts from common defangs.

    :param ioc: String IP Address or URL netloc.
    :rtype: str
    """
    return ioc.replace('[dot]', '.').\
               replace('(dot)', '.').\
               replace('[.]', '.').\
               replace('(', '').\
               replace(')', '').\
               replace(',', '.').\
               replace(' ', '').\
               replace(u'\u30fb', '.')

def refang_url(url):
    """Refang a URL.

    :param url: String URL
    :rtype: str
    """
    # First fix urlparse errors.
    # Fix ipv6 parsing exception.
    if '[.' in url and '[.]' not in url:
        url = url.replace('[.', '[.]')
    if '.]' in url and '[.]' not in url:
        url = url.replace('.]', '[.]')
    if '[dot' in url and '[dot]' not in url:
        url = url.replace('[dot', '[.]')
    if 'dot]' in url and '[dot]' not in url:
        url = url.replace('dot]', '[.]')
    if '[/]' in url:
        url = url.replace('[/]', '/')

    try:
        _ = urlparse(url)
    except ValueError:
        # Last resort on ipv6 fail.
        url = url.replace('[', '').replace(']', '')

    # Since urlparse expects a scheme, make sure one exists.
    if '//' not in url:
        if '__' in url[:7]:
            # Support http__domain.
            url = url.replace('__', '://', 1)
        else:
            # Support no-protocol.
            url = 'http://' + url

    # Refang (/), since it's not entirely in the netloc.
    url = url.replace('(/)', '/')

    # Now use urlparse and continue processing.
    parsed = urlparse(url)

    # Handle URLs with no scheme / obfuscated scheme.
    # Note: ParseResult._replace is a public member, this is safe.
    if parsed.scheme not in ['http', 'https', 'ftp']:
        parsed = parsed._replace(scheme='http')
        url = parsed.geturl().replace('http:///', 'http://')
        parsed = urlparse(url)

    # Remove artifacts from common defangs.
    parsed = parsed._replace(netloc=_refang_common(parsed.netloc))

    # Fix example[.]com, but keep RFC 2732 URLs intact.
    if not _is_ipv6_url(url):
        parsed = parsed._replace(netloc=parsed.netloc.replace('[', '').replace(']', ''))

    return parsed.geturl()

def refang_ipv4(ip_address):
    """Refang an IPv4 address.

    :param ip_address: String IPv4 address.
    :rtype: str
    """
    return _refang_common(ip_address).replace('[', '').replace(']', '')

def main():
    """Run as a commandline utility."""
    parser = argparse.ArgumentParser(
        description="""Advanced Indicator of Compromise (IOC) extractor.
                       If no arguments are specified, the default behavior is
                       to extract all IOCs.""")
    parser.add_argument('--input', type=argparse.FileType('r'),
                        default=sys.stdin, help="default: stdin")
    parser.add_argument('--output', type=argparse.FileType('w'),
                        default=sys.stdout, help="default: stdout")
    parser.add_argument('--extract-ips', action='store_true')
    parser.add_argument('--extract-urls', action='store_true')
    parser.add_argument('--extract-yara-rules', action='store_true')
    parser.add_argument('--extract-hashes', action='store_true')
    parser.add_argument('--refang', action='store_true', help="default: no")
    args = parser.parse_args()

    # By default, extract all
    if not (args.extract_ips or args.extract_urls or args.extract_yara_rules or args.extract_hashes):
        for ioc in extract_iocs(args.input.read(), refang=args.refang):
            args.output.write("{ioc}\r\n".format(ioc=ioc))
    else:
        if args.extract_ips:
            for ioc in extract_ips(args.input.read(), refang=args.refang):
                args.output.write("{ioc}\r\n".format(ioc=ioc))
        if args.extract_urls:
            for ioc in extract_ips(args.input.read(), refang=args.refang):
                args.output.write("{ioc}\r\n".format(ioc=ioc))
        if args.extract_yara_rules:
            for ioc in extract_yara_rules(args.input.read()):
                args.output.write("{ioc}\r\n".format(ioc=ioc))
        if args.extract_hashes:
            for ioc in extract_hashes(args.input.read()):
                args.output.write("{ioc}\r\n".format(ioc=ioc))

if __name__ == "__main__":
    main()
