"""
apooxml - Generate YARA rules for OOXML documents.
author - Aaron Stephens <aaron.stephens@mandiant.com>

Copyright 2021 FireEye, Inc.

Licensed under the Apache License, Version 2.0 (the "License"); you may not
use this file except in compliance with the License. You may obtain a copy
of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
License for the specific language governing permissions and limitations under
the License.
"""


from argparse import ArgumentParser
from datetime import datetime
from hashlib import md5
from io import FileIO
from zipfile import ZipFile


BASE_OOXML_RULE = '''rule {} {{
    meta:
        author = "{}"
        description = "Detects specific {} in OOXML. Generated by apooxml."
        reference_md5 = "{}"

    strings:
        $crc = {{{}}}
        $ext = "{}"
        $ufs = {{{}}}

    condition:
        $ufs at @crc[1] + 8 and $ext at @crc[1] + uint16(@crc[1] + 12) + 16 - {}
}}\n
'''


def select_zip_info(sample: bytes) -> tuple:
    """Print a list of items contained within the ZIP file, along with
    their last modified times, CRC32 checksums, and file sizes. Return
    info on the item selected by the user as a tuple.
    """

    t = []
    w = 0
    z = ZipFile(sample)

    for i in z.infolist():
        if len(i.filename) > w:
            w = len(i.filename)

        t.append((i.filename, datetime(*i.date_time), i.CRC, i.file_size))

    for i in range(len(t)):
        dt = t[i][1].strftime('%Y-%m-%d %H:%M:%S')
        crc = t[i][2].to_bytes(4, 'big').hex()
        print(f'{i + 1: >2}. {t[i][0]: <{w}}  {dt}  {crc}  {t[i][3]}')

    n = input('\nEnter a number corresponding to the desired entry: ')

    print()

    return t[int(n) - 1]


def create_ooxml_rule(author: str, sample: FileIO, name: str = '') -> str:
    """Create YARA rule for given OOXML document."""

    md5_hash = md5(sample.read()).hexdigest()

    file_name, _, crc, size = select_zip_info(sample)

    crc = crc.to_bytes(4, 'little').hex()
    ext = file_name.split('.')[-1]
    ufs = size.to_bytes(4, 'little').hex()

    if not name:
        name = f'ooxml_{ext}_crc_{md5_hash}'

    return BASE_OOXML_RULE.format(
        name, author, ext.upper(), md5_hash, crc, f'.{ext}', ufs, len(ext) + 1
    )


def main(
    sample: str, author: str = 'apooxml', name: str = '', out: str = 'apooxml.yara'
) -> None:
    """Main function for command-line usage."""

    with open(sample, 'rb') as s:
        rule = create_ooxml_rule(author, s, name=name)

    with open(out, 'w') as output:
        output.write(rule)

    print(f'Wrote YARA rule to {out}.')


if __name__ == '__main__':
    parser = ArgumentParser(description='Generate YARA rules for OOXML documents.')

    parser.add_argument('sample', help='OOXML document to generate YARA rule from.')

    parser.add_argument('-a', '--author', default='apooxml', help='YARA rule author.')

    parser.add_argument('-n', '--name', default='', help='YARA rule name.')

    parser.add_argument(
        '-o', '--out', default='apooxml.yara', help='YARA rule file name.'
    )

    main(**vars(parser.parse_args()))
