mb

.gitignore

@@ -0,0 +1,10 @@
+__pycache__/
+.ipynb_checkpoints
+
+.env/
+deployment_files/.*
+deployment_files/client_dir/
+deployment_files/server_dir/
+
+TODO.md
+.venv

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/Zamark.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,15 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="2">
+            <item index="0" class="java.lang.String" itemvalue="concrete-ml" />
+            <item index="1" class="java.lang.String" itemvalue="streamlit" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.12 (Team8)" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/Zamark.iml" filepath="$PROJECT_DIR$/.idea/Zamark.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

app.py

@@ -1,4 +1,254 @@
+
+
 import streamlit as st
+import hashlib
+import uuid
+from streamlit_card import card
+import streamlit.components.v1 as components
+import time
+import json
+
+def generate_mock_hash():
+    return hashlib.sha256(str(time.time()).encode()).hexdigest()
+
+
+from utils import (
+    CLIENT_DIR,
+    CURRENT_DIR,
+    DEPLOYMENT_DIR,
+    KEYS_DIR,
+    INPUT_BROWSER_LIMIT,
+    clean_directory,
+    SERVER_DIR,
+)
+
+from concrete.ml.deployment import FHEModelClient
+
+st.set_page_config(layout="wide")
+
+st.sidebar.title("Contact")
+st.sidebar.info(
+    """
+    - Reda Bellafqira
+    - Mehdi Ben Ghali
+    - Pierre-Elisée Flory
+    - Mohammed Lansari
+    - Thomas Winninger
+    """
+)
+
+st.title("Secure Watermarking Service")
+
+# st.image(
+#     "llm_watermarking.png",
+#     caption="A Watermark for Large Language Models (https://doi.org/10.48550/arXiv.2301.10226)",
+# )
+
+
+def todo():
+    st.warning("Not implemented yet", icon="⚠️")
+
+
+def key_gen_fn(client_id):
+    """
+    Generate keys for a given user. The keys are saved in KEYS_DIR
+
+    !!! needs a model in DEPLOYMENT_DIR as "client.zip" !!!
+    Args:
+        client_id (str): The client_id, retrieved from streamlit
+    """
+    clean_directory()
+
+    client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{client_id}")
+    client.load()
+
+    # Creates the private and evaluation keys on the client side
+    client.generate_private_and_evaluation_keys()
+
+    # Get the serialized evaluation keys
+    serialized_evaluation_keys = client.get_serialized_evaluation_keys()
+    assert isinstance(serialized_evaluation_keys, bytes)
+
+    # Save the evaluation key
+    evaluation_key_path = KEYS_DIR / f"{client_id}/evaluation_key"
+    with evaluation_key_path.open("wb") as f:
+        f.write(serialized_evaluation_keys)
+
+    # show bit of key
+    serialized_evaluation_keys_shorten_hex = serialized_evaluation_keys.hex()[
+        :INPUT_BROWSER_LIMIT
+    ]
+    # shpw len of key
+    # f"{len(serialized_evaluation_keys) / (10**6):.2f} MB"
+    with st.expander("Generated keys"):
+        st.write(f"{len(serialized_evaluation_keys) / (10**6):.2f} MB")
+        st.code(serialized_evaluation_keys_shorten_hex)
+
+    st.success("Keys have been generated!", icon="✅")
+
+
+def gen_trigger_set(client_id, hf_id):
+    # input : random images seeded by client_id
+    # labels : binary array of the id
+    watermark_uuid = uuid.uuid1()
+    hash = hashlib.sha256()
+    hash.update(client_id + str(watermark_uuid))
+    client_seed = hash.digest()
+    hash = hashlib.sha256()
+    hash.update(hf_id + str(watermark_uuid))
+    hf_seed = hash.digest()
+
+    trigger_set_size = 128
+
+    trigger_set_client = [
+        {"input": 1, "label": digit} for digit in encode_id(client_id, trigger_set_size)
+    ]
+
+    todo()
+
+
+def encode_id(ascii_rep, size=128):
+    """Encode a string id to a string of bits
+
+    Args:
+        ascii_rep (_type_): The id string
+        size (_type_): The size of the output bit string
+
+    Returns:
+        _type_: a string of bits
+    """
+    return "".join([format(ord(x), "b").zfill(8) for x in client_id])[:size]
+
+
+def decode_id(binary_rep):
+    """Decode a string of bits to an ascii string
+
+    Args:
+        binary_rep (_type_): the binary string
+
+    Returns:
+        _type_: an ascii string
+    """
+    # Initializing a binary string in the form of
+    # 0 and 1, with base of 2
+    binary_int = int(binary_rep, 2)
+    # Getting the byte number
+    byte_number = binary_int.bit_length() + 7 // 8
+    # Getting an array of bytes
+    binary_array = binary_int.to_bytes(byte_number, "big")
+    # Converting the array into ASCII text
+    ascii_text = binary_array.decode()
+    # Getting the ASCII value
+    return ascii_text
+
+
+def compare_id(client_id, binary_triggert_set_result):
+    """Compares the string id with the labels of the trigger set on the tested API
+
+    Args:
+        client_id (_type_): the ascii string
+        binary_triggert_set_result (_type_): the binary string
+
+    Returns:
+        _type_: _description_
+    """
+    ground_truth = encode_id(client_id, 128)
+
+    correct_bit = 0
+    for true_bit, real_bit in zip(ground_truth, binary_triggert_set_result):
+        if true_bit != real_bit:
+            correct_bit += 1
+
+    return correct_bit / len(binary_triggert_set_result)
+
+
+def watermark(model, trigger_set):
+    """Watermarking function
+
+    Args:
+        model (_type_): The model to watermark
+        trigger_set (_type_): the trigger set
+    """
+    todo()
+
+    model_file_path = SERVER_DIR / "watermarked_model"
+    trigger_set_file_path = SERVER_DIR / "trigger_set"
+
+    # TODO: remove once model correctly watermarked
+    model_file_path.touch()
+    trigger_set_file_path.touch()
+
+    # Once the model is watermarked and dumped to files (model + trigger set), the user can download them
+    with open(model_file_path, "rb") as model_file:
+        st.download_button(
+            label="Download the watermarked file",
+            data=model_file,
+            mime="application/octet-stream",
+        )
+    with open(trigger_set_file_path, "rb") as trigger_set_file:
+        st.download_button(
+            label="Download the triggert set",
+            data=trigger_set_file,
+            mime="application/octet-stream",
+        )
+
+
+st.header("Client Configuration", divider=True)
+
+client_id = st.text_input("Identification string", "team-8-uuid")
+
+if st.button("Generate keys"):
+    key_gen_fn(client_id)
+
+st.header("Model Watermarking", divider=True)
+
+encrypted_model = st.file_uploader("Upload your encrypted model")
+
+if st.button("Start Watermarking"):
+    watermark(None, None)
+
+st.header("Watermarking Verification", divider=True)
+
+
+st.header("Update Blockchain", divider=True)
+
+# Initialize session state to store the block data
+if 'block_data' not in st.session_state:
+    st.session_state.block_data = None
+
+# Button to update the blockchain
+if st.button("Update Blockchain"):
+    previous_hash = generate_mock_hash()
+    timestamp = int(time.time() * 1000)  # Current timestamp in milliseconds
+    watermarked_model_hash = generate_mock_hash()
+    trigger_set_hash = generate_mock_hash()
+
+    # Create the block data structure
+    st.session_state.block_data = {
+        "blockNumber": 42,
+        "previousHash": previous_hash,
+        "timestamp": timestamp,
+        "transactions": [
+            {
+                "type": "Watermarked Model Hash",
+                "hash": watermarked_model_hash
+            },
+            {
+                "type": "Trigger Set Hash",
+                "hash": trigger_set_hash
+            }
+        ]
+    }
+
+    st.success("Blockchain updated successfully!")
+
+# Display the JSON if block_data exists
+if st.session_state.block_data:
+    st.subheader("Latest Block Data (JSON)")
+
+    # Convert the data to a formatted JSON string
+    block_json = json.dumps(st.session_state.block_data, indent=2)
+
+    # Display the JSON
+    st.code(block_json, language='json')
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)

utils.py

@@ -0,0 +1,24 @@
+import os
+import shutil
+from pathlib import Path
+
+INPUT_BROWSER_LIMIT = 380
+
+CURRENT_DIR = Path(__file__).parent
+DEPLOYMENT_DIR = CURRENT_DIR / "deployment_files"
+KEYS_DIR = DEPLOYMENT_DIR / ".fhe_keys"
+CLIENT_DIR = DEPLOYMENT_DIR / "client_dir"
+SERVER_DIR = DEPLOYMENT_DIR / "server_dir"
+
+ALL_DIRS = [KEYS_DIR, CLIENT_DIR, SERVER_DIR]
+
+
+def clean_directory() -> None:
+    """
+    Clear direcgtories
+    """
+    print("Cleaning...\n")
+    for target_dir in ALL_DIRS:
+        if os.path.exists(target_dir) and os.path.isdir(target_dir):
+            shutil.rmtree(target_dir)
+        target_dir.mkdir(exist_ok=True, parents=True)