utils.py

# Copyright (c) 2020, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# For full license text, see the LICENSE file in the repo root
# or https://opensource.org/licenses/BSD-3-Clause

import json
import os
import sys
from hashlib import sha512

import lz4.frame
from Crypto.PublicKey import RSA

from ai_economist.foundation.base.base_env import BaseEnvironment


def save_episode_log(game_object, filepath, compression_level=16):
    """Save a lz4 compressed version of the dense log stored
    in the provided game object"""
    assert isinstance(game_object, BaseEnvironment)
    compression_level = int(compression_level)
    if compression_level < 0:
        compression_level = 0
    elif compression_level > 16:
        compression_level = 16

    with lz4.frame.open(
        filepath, mode="wb", compression_level=compression_level
    ) as log_file:
        log_bytes = bytes(
            json.dumps(
                game_object.previous_episode_dense_log, ensure_ascii=False
            ).encode("utf-8")
        )
        log_file.write(log_bytes)


def load_episode_log(filepath):
    """Load the dense log saved at provided filepath"""
    with lz4.frame.open(filepath, mode="rb") as log_file:
        log_bytes = log_file.read()
    return json.loads(log_bytes)


def verify_activation_code():
    """
    Validate the user's activation code.
    If the activation code is valid, also save it in a text file for future reference.
    If the activation code is invalid, simply exit the program
    """
    path_to_activation_code_dir = os.path.dirname(os.path.abspath(__file__))

    def validate_activation_code(code, msg=b"covid19 code activation"):
        filepath = os.path.abspath(
            os.path.join(
                path_to_activation_code_dir,
                "scenarios/covid19/key_to_check_activation_code_against",
            )
        )
        with open(filepath, "r") as fp:
            key_pair = RSA.import_key(fp.read())

        hashed_msg = int.from_bytes(sha512(msg).digest(), byteorder="big")
        signature = pow(hashed_msg, key_pair.d, key_pair.n)
        try:
            exp_from_code = int(code, 16)
            hashed_msg_from_signature = pow(signature, exp_from_code, key_pair.n)

            return hashed_msg == hashed_msg_from_signature
        except ValueError:
            return False

    activation_code_filename = "activation_code.txt"

    filepath = os.path.join(path_to_activation_code_dir, activation_code_filename)
    if activation_code_filename in os.listdir(path_to_activation_code_dir):
        print("Using the activation code already present in '{}'".format(filepath))
        with open(filepath, "r") as fp:
            activation_code = fp.read()
            fp.close()
        if validate_activation_code(activation_code):
            return  # already activated
        print(
            "The activation code saved in '{}' is incorrect! "
            "Please correct the activation code and try again.".format(filepath)
        )
        sys.exit(0)
    else:
        print(
            "In order to run this simulation, you will need an activation code.\n"
            "Please fill out the form at "
            "https://forms.gle/dJ2gKDBqLDko1g7m7 and we will send you an "
            "activation code to the provided email address.\n"
        )
        num_attempts = 5
        attempt_num = 0
        while attempt_num < num_attempts:
            activation_code = input(
                f"Whenever you are ready, "
                "please enter the activation code: "
                f"(attempt {attempt_num + 1} / {num_attempts})"
            )
            attempt_num += 1
            if validate_activation_code(activation_code):
                print(
                    "Saving the activation code in '{}' for future "
                    "use.".format(filepath)
                )
                with open(
                    os.path.join(path_to_activation_code_dir, activation_code_filename),
                    "w",
                ) as fp:
                    fp.write(activation_code)
                    fp.close()
                return
            print("Incorrect activation code. Please try again.")
        print(
            "You have had {} attempts to provide the activate code. Unfortunately, "
            "none of the activation code(s) you provided could be validated. "
            "Exiting...".format(num_attempts)
        )
        sys.exit(0)