vpnunit/main.py

import sqlite3
import os
from flask import Flask, request, jsonify, render_template, Response
app = Flask(__name__)

class Ex(Exception):
    def __init__(self, code, message):
        self._code = code
        self._message = message
    def getCode(self):
        return self._code
    def __str__(self):
        return self._message

DATABASE='/data/database.sqlite3'
db = sqlite3.connect(DATABASE)
cu = db.cursor()
db.execute('SELECT name FROM sqlite_master')
if len(cu.fetchall()) == 0:
    print('creating database schema...')
    db.execute('CREATE TABLE IF NOT EXISTS user (id INTEGER PRIMARY KEY, name TEXT UNIQUE)')
    db.execute('CREATE TABLE IF NOT EXISTS gateway (id INTEGER PRIMARY KEY, subid INTEGER NOT NULL, name TEXT UNIQUE, user INTEGER NOT NULL REFERENCES user(id))')
    db.commit()
cu.close()
db.close()

os.environ['EASYRSA_PKI'] = '/data/pki'
os.environ['EASYRSA_BATCH'] = '1'
os.environ['PATH'] = os.environ['PATH'] + ':' + os.getcwd() + '/easy-rsa/easyrsa3'

@app.route('/')
def hello_world():
    return 'It works!'

@app.route('/users', methods=['GET'])
def get_users():
    db = sqlite3.connect(DATABASE)
    cu = db.cursor()

    users = []
    for row in cu.execute('SELECT id, name FROM user'):
        users.append({'id': row[0], 'name': row[1]})

    cu.close()
    db.close()
    return jsonify(users)

@app.route('/users', methods=['POST'])
def post_users():
    db = sqlite3.connect(DATABASE)
    cu = db.cursor()

    name = request.json['name']
    print('creating ' + str(name))

    try:
        cu.execute('INSERT INTO user (name) VALUES (?)', (name,))
        return jsonify({'status': 'ok'})
    except sqlite3.Error as e:
        return jsonify({'status': 'error', 'message': str(e)}), 409
    finally:
        db.commit()
        cu.close()
        db.close()

@app.route('/gateways', methods=['GET'])
def get_gateways():
    db = sqlite3.connect(DATABASE)
    cu = db.cursor()

    gateways = []
    for row in cu.execute('SELECT g.id, g.subid, g.name, u.name FROM gateway AS g INNER JOIN user AS u ON u.id = g.user'):
        gateways.append({'id': row[0], 'subid': row[1], 'name': row[2], 'user': row[3]})

    cu.close()
    db.close()
    return jsonify(gateways)

@app.route('/gateways', methods=['POST'])
def post_gateways():
    db = sqlite3.connect(DATABASE)
    cu = db.cursor()

    try:
        name = request.json['name']
        user = request.json['user']

        # TODO sanitize name, it must be a FQDN

        used_gateway_subids = []
        for row in cu.execute('SELECT g.subid, u.id FROM gateway AS g INNER JOIN user AS u ON g.user = u.id WHERE u.name = ?', [str(user,)]):
            used_gateway_subids.append(row[0])
            userid = row[1]

        # search for an empty id for gateway
        for subid in range(0, 15):
            if subid not in used_gateway_subids:
                break

        if subid in used_gateway_subids:
            raise Ex(403, 'exit: maximum number of gateways reached for user')

        cu.execute('INSERT INTO gateway (subid, name, user) VALUES (?, ?, (SELECT id FROM user WHERE name = ?))', [subid, str(name,), str(user,)])

        os.environ['EASYRSA_REQ_CN'] = name

        r = os.system('easyrsa gen-req {} nopass'.format(name))
        if r != 0:
            raise Ex(500, 'exit: {} cannot gen-req'.format(r))
        r = os.system('easyrsa sign-req client {}'.format(name))
        if r != 0:
            raise Ex(500, 'exit: {} cannot sign-req'.format(r))

        ipid = '{:x}{:x}'.format(userid, subid)

        address = '2001:470:c844::' + ipid + '0/64'
        network = '2001:470:c844:' + ipid + '0::/60'

        staticclient = render_template('staticclient', address=address, network=network)
        with open('/data/ovpn/clients/' + name, 'w') as f:
            f.write(staticclient)

        db.commit()
        return jsonify({'status': 'ok'})
    except KeyError as e:
        return jsonify({'status': 'error', 'message': str(e)}), 400
    except sqlite3.Error as e:
        return jsonify({'status': 'error', 'message': str(e)}), 409
    except Ex as e:
        return jsonify({'status': 'error', 'message': str(e)}), e.getCode()
    finally:
        cu.close()
        db.close()

@app.route('/gateway/<fqdn>', methods=['GET'])
def get_gateway(fqdn):
    db = sqlite3.connect(DATABASE)
    cu = db.cursor()

    gateway = dict()
    for row in cu.execute('SELECT g.id, g.subid, g.name, u.name FROM gateway AS g INNER JOIN user AS u ON u.id = g.user'):
        gateway['id'] = row[0]
        gateway['subid'] = row[1]
        gateway['name'] = row[2]
        gateway['user'] = row[3]

    cu.close()
    db.close()

    return jsonify(gateway)

@app.route('/gateway/<fqdn>/config', methods=['GET'])
def get_gateway_config(fqdn):
    # TODO sanity check FQDN
    # WARNING: maybe you want to do more than a simple sanity check,
    # or you have to trust the user of these API
    # eg: he could retrieve the CA.key !!!

    with open(os.environ['EASYRSA_PKI'] + '/issued/' + fqdn + '.crt') as f:
        cert = f.read()
    with open(os.environ['EASYRSA_PKI'] + '/private/' + fqdn + '.key') as f:
        key = f.read()
    with open(os.environ['EASYRSA_PKI'] + '/ca.crt') as f:
        ca = f.read()

    return Response(render_template('config.ovpn', ca=ca, cert=cert, key=key), mimetype='text/plain')

if __name__ == '__main__':
    app.run(host="::", port=5000, debug=True)