generate.py

#!/usr/bin/env python3

import copy
import glob
import jinja2
import jinja2.ext
import os
import shutil
import subprocess
import yaml

def run_subprocess(command, outfilename=None, working_dir='.', expected_returncode=0, input=None, ignore_returncode=False):
    print(working_dir + " $ " + " ".join(command) +" > "+outfilename)

    result = subprocess.run(
            command,
            input=input,
            stdout=(open(outfilename, "w") if outfilename!=None else subprocess.PIPE),
            cwd=working_dir,
        )

    if not(ignore_returncode) and (result.returncode != expected_returncode):
        if outfilename == None:
            print(result.stdout.decode('utf-8'))
        assert False, "Got unexpected return code {}".format(result.returncode)
    return result.stderr.decode('utf-8') if result.stderr else None


# Function replicating openssl makefile target of the same name to generate ASN1 structs
# The function assumes perl to be present as well as an openssl master branch checked out.
def generate_crypto_objects():
    OSSL_SRC_DIR = 'openssl'
    OSSL_OBJ_GEN_DIR = os.path.join(OSSL_SRC_DIR, 'crypto', 'objects')
    # check presence of OpenSSL (script files):
    if not os.path.exists(OSSL_OBJ_GEN_DIR):
        print("Error: OSSL install not present (%s missing). Exiting." % (OSSL_OBJ_GEN_DIR))
        exit(1)
    # now run the generator scripts:
    cmd = ['perl', os.path.join(OSSL_OBJ_GEN_DIR, 'objects.pl'), '-n' , os.path.join('oqsprov', 'oqs_objects.txt') , os.path.join('oqsprov', 'oqs_obj_mac.num')  ]
    run_subprocess(cmd, os.path.join('oqsprov', 'oqs_obj_mac.new'))
    os.rename(os.path.join('oqsprov', 'oqs_obj_mac.new'), os.path.join('oqsprov', 'oqs_obj_mac.num'))
    cmd = ['perl', os.path.join(OSSL_OBJ_GEN_DIR, 'objects.pl'), os.path.join('oqsprov', 'oqs_objects.txt'), os.path.join('oqsprov', 'oqs_obj_mac.num') ]
    run_subprocess(cmd, os.path.join('oqsprov', 'oqs_obj_mac.h'))
    cmd = ['perl', os.path.join(OSSL_OBJ_GEN_DIR, 'obj_dat.pl'), os.path.join('oqsprov', 'oqs_obj_mac.h')]
    run_subprocess(cmd, os.path.join('oqsprov', 'oqs_obj_dat.h'))
    cmd = ['perl', os.path.join(OSSL_OBJ_GEN_DIR, 'objxref.pl'), os.path.join('oqsprov', 'oqs_obj_mac.num'), os.path.join('oqsprov', 'oqs_obj_xref.txt')]
    run_subprocess(cmd, os.path.join('oqsprov', 'oqs_obj_xref.h'))
    
# For list.append in Jinja templates
Jinja2 = jinja2.Environment(loader=jinja2.FileSystemLoader(searchpath="."),extensions=['jinja2.ext.do'])

def file_get_contents(filename, encoding=None):
    with open(filename, mode='r', encoding=encoding) as fh:
        return fh.read()

def file_put_contents(filename, s, encoding=None):
    with open(filename, mode='w', encoding=encoding) as fh:
        fh.write(s)

def populate(filename, config, delimiter, overwrite=False):
    fragments = glob.glob(os.path.join('oqs-template', filename, '*.fragment'))
    if overwrite == True:
        source_file = os.path.join('oqs-template', filename, os.path.basename(filename)+ '.base')
        contents = file_get_contents(source_file)
    else:
        contents = file_get_contents(filename)
    for fragment in fragments:
        identifier = os.path.splitext(os.path.basename(fragment))[0]
        identifier_start = '{} OQS_TEMPLATE_FRAGMENT_{}_START'.format(delimiter, identifier.upper())
        identifier_end = '{} OQS_TEMPLATE_FRAGMENT_{}_END'.format(delimiter, identifier.upper())
        preamble = contents[:contents.find(identifier_start)]
        postamble = contents[contents.find(identifier_end):]
        if overwrite == True:
            contents = preamble + Jinja2.get_template(fragment).render({'config': config}) + postamble.replace(identifier_end + '\n', '')
        else:
            contents = preamble + identifier_start + Jinja2.get_template(fragment).render({'config': config}) + postamble
    file_put_contents(filename, contents)

def load_config():
    config = file_get_contents(os.path.join('oqs-template', 'generate.yml'), encoding='utf-8')
    config_extras = file_get_contents(os.path.join('oqs-template', 'generate-extras.yml'), encoding='utf-8')
    config = yaml.safe_load(config)
    config_extras = yaml.safe_load(config_extras)
    for sig in config['sigs']:
        sig['variants'] = [variant for variant in sig['variants'] if variant['enable']]

    # remove KEMs without NID (old stuff)
    newkems = []
    for kem in config['kems']:
        if 'nid' in kem:
           newkems.append(kem)
    config['kems']=newkems

    for kem in config['kems']:
        if kem['name_group'] in config_extras['kem-extras']:
            kem.update(config_extras['kem-extras'][kem['name_group']])
        try:
            for extra_nid_current in kem['extra_nids']['current']:
                if 'hybrid_group' in extra_nid_current and extra_nid_current['hybrid_group'] in ["x25519", "x448"]:
                    extra_hyb_nid = extra_nid_current['nid']
                    if 'nid_ecx_hybrid' in kem:
                        print("Warning, duplicate nid_ecx_hybrid for",
                              kem['name_group'], ":", extra_hyb_nid, "in generate.yml,",
                              kem['nid_ecx_hybrid'], "in generate_extras.yml, using generate.yml entry.")
                    kem['nid_ecx_hybrid'] = extra_hyb_nid
                    break
        except:
            pass
    return config

config = load_config()

populate('test/oqs_test_signatures.c', config, '/////')
populate('test/oqs_test_groups.c', config, '/////')
populate('oqsprov/oqs_prov.h', config, '/////')
populate('oqsprov/oqsprov.c', config, '/////')
populate('oqsprov/oqsprov_groups.c', config, '/////')
populate('oqsprov/oqs_kmgmt.c', config, '/////')
populate('oqsprov/oqs_sig.c', config, '/////')
populate('oqsprov/oqs_objects.txt', config, '#####')
populate('oqsprov/oqs_obj_xref.txt', config, '#####')

generate_crypto_objects()