#!/usr/bin/env python3

"""
(C) 2016 Jack Lloyd
(C) 2017 Fabian Weissberg, Rohde & Schwarz Cybersecurity

Botan is released under the Simplified BSD License (see license.txt)
"""

import sys
import datetime
import re
from collections import defaultdict

def format_oid(oid):
    #return '"' + oid + '"'
    return "{" + oid.replace('.', ', ') + '}'

def format_map(m, for_oid = False):
    s = ''
    for k in sorted(m.keys()):
        v = m[k]

        if len(s) > 0:
            s += '      '

        if for_oid:
            s += '{"%s", OID(%s)},\n' % (k,format_oid(v))
        else:
            s += '{"%s", "%s"},\n' % (k,v)

    s = s[:-2] # chomp last two chars

    return s


def format_as_map(oid2str, str2oid):
   return """/*
* OID maps
*
* This file was automatically generated by %s on %s
*
* All manual edits to this file will be lost. Edit the script
* then regenerate this source file.
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/internal/oid_map.h>
#include <unordered_map>

namespace Botan {

std::unordered_map<std::string, std::string> OID_Map::load_oid2str_map() {
   return std::unordered_map<std::string, std::string>{

      %s};
}

std::unordered_map<std::string, OID> OID_Map::load_str2oid_map() {
   return std::unordered_map<std::string, OID>{

      %s};
}

}  // namespace Botan""" % (
    sys.argv[0],
    datetime.date.today().strftime("%Y-%m-%d"),
    format_map(oid2str),
    format_map(str2oid, True))

def format_dn_ub_map(dn_ub, oid2str):
    s = ''
    for k in sorted(dn_ub.keys()):
        v = dn_ub[k]

        expr = "   {OID({%s}), %s}, " % (k.replace('.', ', '), v)
        s += expr
        s += ' '*(32 - len(expr))
        s += ' // %s\n' % (oid2str[k])

    # delete last ',' and \n
    idx = s.rfind(',')
    if idx != -1:
        s = s[:idx] + ' ' + s[idx+1:-1]

    return s


def format_dn_ub_as_map(dn_ub, oid2str):
    return """/*
* DN_UB maps: Upper bounds on the length of DN strings
*
* This file was automatically generated by %s on %s
*
* All manual edits to this file will be lost. Edit the script
* then regenerate this source file.
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/pkix_types.h>

#include <botan/asn1_obj.h>
#include <map>

namespace Botan {

namespace {

/**
 * Upper bounds for the length of distinguished name fields as given in RFC 5280, Appendix A.
 * Only OIDS recognized by botan are considered, so far.
 * Maps OID string representations instead of human readable strings in order
 * to avoid an additional lookup.
 */
const std::map<OID, size_t> DN_UB = {
%s
};

}  // namespace

//static
size_t X509_DN::lookup_ub(const OID& oid) {
   auto ub_entry = DN_UB.find(oid);
   if(ub_entry != DN_UB.end()) {
      return ub_entry->second;
   } else {
      return 0;
   }
}

}  // namespace Botan""" % (sys.argv[0], datetime.date.today().strftime("%Y-%m-%d"),
       format_dn_ub_map(dn_ub,oid2str))


def format_set_map(m):
    s = ''
    for k in sorted(m.keys()):
        v = m[k]

        if len(s) > 0:
            s += '   '

        s += '{ "%s", {' % k
        for pad in v:
            s += '"%s", ' % pad
        if len(v) != 0:
            s = s[:-2]
        s += '} },\n'
    s = s[:-1]
    return s


def main(args = None):
    """ Print header files (oids.cpp, dn_ub.cpp) depending on the first argument and on src/build-data/oids.txt

        Choose 'oids' to print oids.cpp, needs to be written to src/lib/asn1/oids.cpp
        Choose 'dn_ub' to print dn_ub.cpp, needs to be written to src/lib/x509/X509_dn_ub.cpp
    """
    if args is None:
        args = sys.argv
    if len(args) < 2:
        print("Use either 'oids' or 'dn_ub' as first argument")
        return 1

    oid_lines = open('./src/build-data/oids.txt').readlines()

    oid_re = re.compile(r"^([0-9][0-9.]+) += +([A-Za-z0-9_\./\(\), -]+)(?: = )?([0-9]+)?$")
    hdr_re = re.compile(r"^\[([a-z0-9_]+)\]$")

    oid2str = {}
    str2oid = {}
    dn_ub = {}
    cur_hdr = None

    for line in oid_lines:
        line = line.strip()
        if len(line) == 0:
            continue

        if line[0] == '#':
            continue

        match = hdr_re.match(line)
        if match is not None:
            cur_hdr = match.group(1)
            continue

        match = oid_re.match(line)
        if match is None:
            raise Exception(line)

        oid = match.group(1)
        nam = match.group(2)

        if oid in str2oid:
            print("Duplicated OID", oid, name, oid2str[oid])
            sys.exit() # hard error
        else:
            oid2str[oid] = nam

        # parse upper bounds for DNs
        if cur_hdr == "dn":
            if match.lastindex < 3:
                raise Exception("Could not find an upper bound for DN " + match.group(1))
            dn_ub[oid] = match.group(3)

        if nam in str2oid:
            #str2oid[nam] = oid
            pass
        else:
            str2oid[nam] = oid

    if args[1] == "oids":
        print(format_as_map(oid2str, str2oid))
    elif args[1] == "dn_ub":
        print(format_dn_ub_as_map(dn_ub,oid2str))
    else:
        print("Unknown command: try oids or dn_ub")

    return 0

if __name__ == '__main__':
    sys.exit(main())
