<commit_before>import argparse
import base64
import json
import struct

from flask import Flask, request
from pyscarab.scarab import EncryptedArray, PublicKey
from utils import binary
from store import Store


parser = argparse.ArgumentParser(description="Start a Blindstore server.")
parser.add_argument('-d', '--debug', action='store_true',
                    help="enable Flask debug mode. DO NOT use in production.")
args = parser.parse_args()

NUM_RECORDS = 5
RECORD_SIZE = 64

def _bytes_to_int(byte_string):
    struct.unpack('<L', byte_string)

app = Flask(__name__)

store = Store()

@app.route('/db_size')
def get_db_size():
    return json.dumps({'num_records': NUM_RECORDS, 'record_size': RECORD_SIZE}), \
        200, {'Content-Type': 'text/json'}


@app.route('/retrieve', methods=['POST'])
def retrieve():
    print("Starting retrieve call...")
    public_key = PublicKey(request.form['PUBLIC_KEY'])

    enc_index = EncryptedArray(store.record_size, public_key, request.form['ENC_INDEX'])
    enc_data = store.retrieve(enc_index, public_key)
    s_bits = [str(b) for b in enc_data]
    obj = json.dumps(s_bits)

    return obj


@app.route('/set', methods=['POST'])
def set():
    try:
        index = int(request.form['INDEX'])
        data = int.from_bytes(base64.b64decode(request.form['DATA']), 'big')

        store.set(index, binary(data, store.record_size))
        return '', 200
    except Exception as e:
        #import traceback
        #traceback.print_last()
        print(e)


if __name__ == '__main__':
    app.run(debug=args.debug)
<commit_msg>Make /db_size return the actual database size<commit_after>import argparse
import base64
import json
import struct

from flask import Flask, request
from pyscarab.scarab import EncryptedArray, PublicKey
from utils import binary
from store import Store


parser = argparse.ArgumentParser(description="Start a Blindstore server.")
parser.add_argument('-d', '--debug', action='store_true',
                    help="enable Flask debug mode. DO NOT use in production.")
args = parser.parse_args()

def _bytes_to_int(byte_string):
    struct.unpack('<L', byte_string)

app = Flask(__name__)

store = Store()

@app.route('/db_size')
def get_db_size():
    data = {'num_records': store.record_count, 'record_size': store.record_size}
    return json.dumps(data), 200, {'Content-Type': 'text/json'}


@app.route('/retrieve', methods=['POST'])
def retrieve():
    print("Starting retrieve call...")
    public_key = PublicKey(request.form['PUBLIC_KEY'])

    enc_index = EncryptedArray(store.record_size, public_key, request.form['ENC_INDEX'])
    enc_data = store.retrieve(enc_index, public_key)
    s_bits = [str(b) for b in enc_data]
    obj = json.dumps(s_bits)

    return obj


@app.route('/set', methods=['POST'])
def set():
    try:
        index = int(request.form['INDEX'])
        data = int.from_bytes(base64.b64decode(request.form['DATA']), 'big')

        store.set(index, binary(data, store.record_size))
        return '', 200
    except Exception as e:
        #import traceback
        #traceback.print_last()
        print(e)


if __name__ == '__main__':
    app.run(debug=args.debug)

{'fraction_non_alphanumeric': 0.13646391427670973, 'fraction_numerical': 0.008509297195083518, 'mean_word_length': 3.8981481481481484, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3647784', 'n_tokens_mistral': 1054, 'n_tokens_neox': 1011, 'n_words': 246}