poc for implementation of polymorphic pseudonyms

gistfile1.rs

use libpep::*;
use libpep::simple::*;
use rand_core::OsRng;
use std::{fmt::Write, num::ParseIntError};
use std::cell::RefCell;
use std::collections::HashMap;
use std::ops::Mul;
use std::rc::Rc;

pub fn decode_hex(s: &str) -> Result<Vec<u8>, ParseIntError> {
    (0..s.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&s[i..i + 2], 16))
        .collect()
}

pub fn encode_hex(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for &b in bytes {
        write!(&mut s, "{:02x}", b).unwrap();
    }
    s
}

/// Keypair is a simple pub/secret key pair
#[derive(Clone)]
struct KeyPair {
    pub pk: GroupElement,
    pub sk: ScalarNonZero,
}

impl KeyPair {
    fn new() -> Self {
        let mut rng = OsRng;
        let (pk, sk) = generate_global_keys(&mut rng);
        Self {
            pk,
            sk,
        }
    }
}

// We need both a K and S to rekey/shuffle. However, we have split these numbers into two parts. One part is stored
// in the transcryptor, the other part is stored in the access manager. This is done to prevent the transcryptor
// from being able to decrypt the data.
// So we have a factorpool for the access_manager, and a factorpool for the transcryptor. To get the full K and S
// we need to combine the two factors:
//
//    K = k_am * k_t
//    S = s_am * s_t
//
//  Only when we rekey/shuffle (rks) with these two factors, we can decrypt the data with the private key of
// that destination.
#[derive(Clone)]
struct Factor {
    k: ScalarNonZero,
    s: ScalarNonZero,
}

/// A factor pool holds a list of factors (K,S numbers) for each destination.
struct FactorPool {
    factors: HashMap<String, Factor>,
}

impl FactorPool {
    pub fn new() -> Self {
        Self {
            factors: HashMap::new(),
        }
    }

    pub fn get(&self, id: &str) -> Option<&Factor> {
        self.factors.get(id)
    }

    pub fn set(&mut self, id: &str, k: ScalarNonZero) {
        let mut rng = OsRng;

        let s = ScalarNonZero::random(&mut rng);
        let factor = Factor{k, s};

        self.factors.insert(id.to_string(), factor);
    }
}

// Key server holds all keys. The global key, and the numbers for each factor. These factors are allowed outside
// the keyserver.
struct KeyServer {
    /// Global keypair for generating things
    global: KeyPair,
    /// Pools to store factor elements inside
    pool_am: Rc<RefCell<FactorPool>>,
    pool_t: Rc<RefCell<FactorPool>>,
}

impl KeyServer {
    pub fn new(pool_am: Rc<RefCell<FactorPool>>, pool_t: Rc<RefCell<FactorPool>>) -> Self {
        Self {
            global: KeyPair::new(),
            pool_am,
            pool_t
        }
    }

    pub fn global_public_key(&self) -> GroupElement {
        return self.global.pk.clone();
    }

    // Don't use this
    pub fn global_secret_key(&self) -> ScalarNonZero {
        println!("This secret key is only retrieved for testing purposes. It should not be used in real life.");
        return self.global.sk.clone();
    }

    // This is defined section 2.3.1 of the whitepaper. It will return Xa, which is generated from Ka, which is
    // generated from Ka@am and Ka@t. The last two are stored in the keyserver but should in real life be sent
    // to the access manager and transcryptor respectively.
    pub fn generate_factor(&mut self, id: &str) -> ScalarNonZero {
        let mut rng = OsRng;

        // Generate Ka@am and Ka@t
        let ka_am = ScalarNonZero::random(&mut rng);
        let ka_t = ScalarNonZero::random(&mut rng);
        let ka = ka_am.mul(&ka_t);

        self.pool_am.borrow_mut().set(id, ka_am.clone());
        self.pool_t.borrow_mut().set(id, ka_t.clone());

        // Xa is the private key for "a"
        let xa = ka.mul(&self.global.sk);
        xa
    }
}

/// Access manager does not do anything yet.
struct AccessManager {
    // key_server: Rc<RefCell<KeyServer>>,
    pool: Rc<RefCell<FactorPool>>,
}

impl AccessManager {
    pub fn new(pool: Rc<RefCell<FactorPool>>) -> Self {
        Self {
            pool,
        }
    }

    fn get_factor(&self, id: &str) -> Factor {
        if let Some(factor) = self.pool.borrow().get(id) {
            return factor.clone();
        }

        panic!("No factor found");
    }

    fn rks(&self, id: &str, data: ElGamal) -> ElGamal {
        let ks = self.get_factor(id);
        rks(&data, &ks.k, &ks.s)
    }

    fn rekey(&self, data: &ElGamal, id: &str) -> ElGamal {
        let factor = self.get_factor(id);
        rekey(&data, &factor.k)
    }
}

struct TransCryptor {
    pool: Rc<RefCell<FactorPool>>,
}

impl TransCryptor {
    fn new(pool: Rc<RefCell<FactorPool>>) -> Self {
        Self {
            pool,
        }
    }

    fn get_factor(&self, id: &str) -> Factor {
        if let Some(factor) = self.pool.borrow().get(id) {
            return factor.clone();
        }

        panic!("No factor found");
    }

    fn rks(&self, id: &str, data: ElGamal) -> ElGamal {
        let factor = self.get_factor(id);
        rks(&data, &factor.k, &factor.s)
    }

    fn rekey(&self, data: &ElGamal, id: &str) -> ElGamal {
        let factor = self.get_factor(id);
        rekey(&data, &factor.k)
    }
}

struct StorageFacility {
    secret: ScalarNonZero,
    data: HashMap<String, Vec<ElGamal>>,
}

fn key_to_string(key: &GroupElement) -> String {
    encode_hex(&key.encode())
}

impl StorageFacility {
    fn new(key_server: Rc<RefCell<KeyServer>>) -> Self {
        Self {
            secret: key_server.borrow_mut().generate_factor("SF"),
            data: HashMap::new(),
        }
    }

    fn store(&mut self, pid: ElGamal, data: ElGamal) {
        let key = key_to_string(&decrypt(&pid, &self.secret));
        println!(">> Storing data for key '{}'", key);

        if !self.data.contains_key(&key) {
            self.data.insert(key.clone(), vec![data]);
        } else {
            self.data.get_mut(&key).expect("no key").push(data);
        }
    }

    fn get(&self, ppid_sf: ElGamal) -> Option<&Vec<ElGamal>> {
        let key = key_to_string(&decrypt(&ppid_sf, &self.secret));
        println!("<< Retrieving data for key '{}'", key);

        // Note that data is encrypted with the global public key, and thus the global secret can decrypt it.

        // the data we get from the storage facility should be rerandomized so it will always return different
        // (encrypted) data, but can always we decrypted to the same plaintext. We already know this works so we
        // don't do this here at this point.
        self.data.get(&key)
    }
}

fn main() {
    let mut rng = OsRng;

    //
    // Global setup
    //

    let pool_am = Rc::new(RefCell::new(FactorPool::new()));
    let pool_t = Rc::new(RefCell::new(FactorPool::new()));
    let key_server = Rc::new(RefCell::new(KeyServer::new(pool_am.clone(), pool_t.clone())));
    let access_manager = AccessManager::new(pool_am.clone());
    let transcryptor = TransCryptor::new(pool_t.clone());
    let mut storage_facility = StorageFacility::new(key_server.clone());

    // The key server has generated a keypair which we can use the public key. All systems are allowed to view this.
    println!("Some global information:");
    println!("  Public global key: {}", encode_hex(&key_server.borrow().global_public_key().encode()));

    println!();
    println!("================================================================================");
    println!();

    //
    // The "smartwatch". A device that is able to communicate with the SF to send over data for the given user (A).
    //

    println!();
    println!("Work done in 'the watch'");

    // Generate pseudonym on a system for user "A". This could be a BSN or anything else that identifies this user. The system
    // does not know the BSN directly, but only the pseudonym. The pseudonym is generated from the BSN and the global public key.
    // We assume that somehow the GEP is sent to the watch.
    let pid_a = "bsn-of-a";
    let ppid_a = generate_pseudonym(pid_a, &key_server.borrow().global_public_key(), &mut rng);
    println!("  PPID(a): {}", encode_hex(&ppid_a.encode()));

    // Create random data that we want to send to the storage facility
    let aes_key1 = GroupElement::random(&mut rng);
    let enc_data = encrypt(&aes_key1, &key_server.borrow().global_public_key(), &mut rng);
    println!("  AES key1: {}", encode_hex(&aes_key1.encode()));

    let ppid_a = rerandomize_local(&ppid_a, &mut rng);
    let ppid_a_sf = transform_for_dest(ppid_a.clone(), "SF", &transcryptor, &access_manager);
    storage_facility.store(ppid_a_sf, enc_data);

    // We send more information (another AES key) to SF. It should be stored in the same LEP as the first key. SF CANNOT decrypt
    // the data, but it can decrypt the pseudonym.
    let aes_key2 = GroupElement::random(&mut rng);
    let enc_data = encrypt(&aes_key2, &key_server.borrow().global_public_key(), &mut rng);
    println!("  AES key2: {}", encode_hex(&aes_key2.encode()));

    let ppid_a = rerandomize_local(&ppid_a, &mut rng);
    let ppid_a_sf = transform_for_dest(ppid_a, "SF", &transcryptor, &access_manager);
    storage_facility.store(ppid_a_sf.clone(), enc_data);

    println!();
    println!("================================================================================");
    println!();

    println!("At the doctor (DOC) now...");
    println!();


    // Try and decode stuff that we can get from the transcryptor. Normally, the access manager will be part of this system
    // because it needs to check if we are allowed to access this data. But for this proof-of-concept we will just assume
    // that we are allowed to access this data.


    let doc_secret = key_server.borrow_mut().generate_factor("DOC");

    // Let's create a doctor (destination: DOC), that we can send the data to. The doctor knows the BSN of the patient
    let ppid_a = generate_pseudonym("bsn-of-a", &key_server.borrow().global_public_key(), &mut rng);
    let ppid_a_sf = transform_for_dest(ppid_a.clone(), "SF", &transcryptor, &access_manager);
    let storage_data = storage_facility.get(ppid_a_sf).unwrap();
    println!("  Data from SF: {} items", storage_data.len());

    // Here we fetch the data straight from the storage facility. This is pointless, as the data coming from the storage
    // can only be decrypted by the global secret. Instead, we should proxy this data through the transcryptor to rerandomize
    // and rekey the data to the doctor.
    let data_doc = storage_data.get(0).unwrap().clone();
    let data_doc = rerandomize(&data_doc, &ScalarNonZero::random(&mut rng));
    let data_doc = transcryptor.rekey(&data_doc, "DOC");
    let data_doc = access_manager.rekey(&data_doc, "DOC");

    // we've got the secret data. Let's try and decrypt it with the global secret first. It should fail since the data is rekeyed only for the doctor.
    let data = decrypt(&data_doc.clone(), &key_server.borrow().global_secret_key());
    println!("  Decrypted by global secret: {}", encode_hex(&data.encode()));

    // Then we try to decrypt it with the doctor secret. This should work.
    let data = decrypt(&data_doc.clone(), &doc_secret);
    println!("  Decrypted by doctor secret: {}", encode_hex(&data.encode()));
}

/// Transforms ppid (polymorphic pseudonym id) into a factor that can be used by the destination.
fn transform_for_dest(ppid: ElGamal, dest: &str, transcryptor: &TransCryptor, access_manager: &AccessManager) -> ElGamal {
    let mut rng = OsRng;

    let ppid = rerandomize_local(&ppid, &mut rng);

    // Let both the transcryptor and access manager rks the data, as both are needed in order for dest to decrypt the data.
    let ppid_dest = transcryptor.rks(dest, ppid);
    let ppid_dest = access_manager.rks(dest, ppid_dest);

    // This is the data we need to send to the dest.
    println!("! epid(A)@{}: {}", dest, encode_hex(&ppid_dest.encode()));

    ppid_dest
}