Skip to content

Commit

Permalink
Model verification added
Browse files Browse the repository at this point in the history
  • Loading branch information
@yugjindal22
yugjindal22 committed Nov 17, 2024
1 parent 5c91013 commit f5b97e8
Show file tree
Hide file tree
Showing 5 changed files with 178 additions and 145 deletions.
152 changes: 77 additions & 75 deletions frontend/src/App.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,44 +265,83 @@ function App() {
throw error; // Re-throw to maintain error chain
}
};

async function fetchPublicVerificationKey() {
// Fetch the PEM formatted public key from the server
const response = await fetch('http://localhost:8000/api/public-verification-key');
if (!response.ok) {
try {
const response = await fetch('http://localhost:8000/api/public-verification-key');
if (!response.ok) {
throw new Error("Failed to fetch public key.");
}
const data = await response.json();
console.log("Received public key data:", data); // Debug log
return await importPublicKey(data.publicKey);
} catch (error) {
console.error("Error fetching public key:", error);
throw error;
}
}

const publicKeyPem = await response.text(); // Assuming the public key is returned as PEM format
async function importPublicKey(pemKey) {
try {
if (!pemKey) {
throw new Error("PEM key is undefined or null");
}

// Import the PEM formatted public key into a CryptoKey object
const publicKey = await importPublicKey(publicKeyPem);
// Clean the PEM key
const pemContents = pemKey
.replace('-----BEGIN PUBLIC KEY-----', '')
.replace('-----END PUBLIC KEY-----', '')
.replace(/[\r\n]+/g, '')
.trim();

return publicKey; // Return the imported CryptoKey object
}
console.log("Cleaned PEM contents length:", pemContents.length); // Debug log

async function importPublicKey(pemKey) {
// Remove the "BEGIN" and "END" parts of the PEM string and decode the base64
const keyParts = pemKey.replace(/-----BEGIN PUBLIC KEY-----|-----END PUBLIC KEY-----|\n/g, '');
const keyBytes = Uint8Array.from(atob(keyParts), c => c.charCodeAt(0));
// Convert from base64 to binary
const binaryString = window.atob(pemContents);
const bytes = new Uint8Array(binaryString.length);
for (let i = 0; i < binaryString.length; i++) {
bytes[i] = binaryString.charCodeAt(i);
}

try {
const importedKey = await crypto.subtle.importKey(
"spki", // Format for public keys
keyBytes,
{
name: "RSA-OAEP", // Algorithm name (or you can use another algorithm, depending on your use case)
hash: { name: "SHA-256" }
},
true, // Extractable
["verify"] // The operations the key will be used for
);
return importedKey;
// Import the key
return await window.crypto.subtle.importKey(
'spki',
bytes.buffer,
{
name: 'RSASSA-PKCS1-v1_5',
hash: { name: 'SHA-256' },
},
true,
['verify']
);
} catch (error) {
console.error("Error importing public key:", error);
throw new Error("Failed to import public key.");
console.error("Error importing public key:", error);
console.error("PEM key received:", pemKey);
throw error;
}
}
}

async function verifySignedHash(originalHash, signedHash, publicKey) {
try {
// Convert the signed hash from base64
const signatureBytes = new Uint8Array(atob(signedHash).split('').map(c => c.charCodeAt(0)));

// Convert the original hash from hex to ArrayBuffer
const messageBytes = new Uint8Array(originalHash.match(/.{1,2}/g).map(byte => parseInt(byte, 16)));

return await window.crypto.subtle.verify(
{
name: 'RSASSA-PKCS1-v1_5',
},
publicKey,
signatureBytes,
messageBytes
);
} catch (error) {
console.error("Error verifying hash:", error);
return false;
}
}

// Function to decrypt the model data using the AES key and IV
const decryptWithAes = async (aesKey, encryptedModelBase64, ivBase64) => {
Expand All @@ -325,19 +364,19 @@ async function importPublicKey(pemKey) {
aesKey,
encryptedModelBuffer
);

// Step 3: Generate hash of the decrypted model
const modelHash = await generateModelHash(decryptedModelBuffer);
console.log("Model hash:", modelHash);
console.log("Signed hash:", signedHash);
const publicVerificationKey = await fetchPublicVerificationKey();
if(await verifySignedHash(modelHash, signedHash, publicVerificationKey)){
const modelHash = await generateModelHash(decryptedModelBuffer);
console.log("Frontend Model hash:", modelHash);
console.log("Frontend Signed Model hash:", signedHash);
const publicVerificationKey = await fetchPublicVerificationKey();
if (await verifySignedHash(modelHash, signedHash, publicVerificationKey)) {

console.log("Model verified successfully, loading model...");
}else{
console.error("Model verification failed, aborting decryption.");
console.log("Model verified successfully, loading model...");
} else {
console.error("Model verification failed, aborting decryption.");

}
}

// Don't try to decode the binary data as text
console.log("Model decrypted successfully! Size:", decryptedModelBuffer.byteLength);
Expand All @@ -358,7 +397,7 @@ async function importPublicKey(pemKey) {
throw new Error("Failed to generate model hash.");
}
}

function arrayBufferToHex(buffer) {
const byteArray = new Uint8Array(buffer);
let hexString = "";
Expand All @@ -368,44 +407,7 @@ async function importPublicKey(pemKey) {
}
return hexString;
}
async function verifySignedHash(originalHash, signedHash, publicKey) {
// Convert the signed hash from Base64 to Uint8Array
const signedHashBuffer = new Uint8Array(atob(signedHash).split("").map(c => c.charCodeAt(0)));

// Convert the original hash from hex string to ArrayBuffer
const originalHashBuffer = hexStringToArrayBuffer(originalHash);

// Use the SubtleCrypto API to verify the signature
const isValid = await crypto.subtle.verify(
{
name: "RSASSA-PKCS1-v1_5", // RSA algorithm for signature verification
},
publicKey, // Public key to verify with
signedHashBuffer, // The signed hash (signature)
originalHashBuffer // The original hash that was signed
);

return isValid; // Returns true if the signature is valid, otherwise false
}

// Utility function to convert a hex string to ArrayBuffer
function hexStringToArrayBuffer(hex) {
const length = hex.length / 2;
const arrayBuffer = new ArrayBuffer(length);
const uint8Array = new Uint8Array(arrayBuffer);

for (let i = 0; i < length; i++) {
uint8Array[i] = parseInt(hex.substr(i * 2, 2), 16);
}
return arrayBuffer;
}
function hexStringToArrayBuffer(hexString) {
const bytes = new Uint8Array(hexString.length / 2);
for (let i = 0; i < hexString.length; i += 2) {
bytes[i / 2] = parseInt(hexString.substr(i, 2), 16);
}
return bytes.buffer;
}
// Utility functions for data conversion between ArrayBuffer and Base64
const arrayBufferToBase64 = (buffer) => {
let binary = "";
Expand Down
54 changes: 27 additions & 27 deletions mern-backend/controllers/modelController.js
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
const {fetchEncryptedFilesFromS3} = require("../utils/s3utils");
const { fetchEncryptedFilesFromS3 } = require("../utils/s3utils");
const { encryptModel } = require("../utils/encryptionUtils.js");
const { generateModelHash, signModelHash } = require("../utils/hashUtils.js");
const fs = require("fs");
const { generateModelHash, signModelHash } = require("../utils/hashUtils.js");
const fs = require("fs");
const path = require("path");
exports.getAllEncryptedModels = async (req, res, next) => {
try {
const modelKey = "antispofing.onnx";
const modelKey = "antispofing.onnx";

const { modelFile } = await fetchEncryptedFilesFromS3(
modelKey
Expand All @@ -14,35 +14,35 @@ exports.getAllEncryptedModels = async (req, res, next) => {
console.log("Decrypted model file:", modelFile);

const publicKeyBase64 = req.body.publicKey;
if (!publicKeyBase64) {
return res.status(400).json({ message: "Public key is required" });
}
if (!publicKeyBase64) {
return res.status(400).json({ message: "Public key is required" });
}

const { encryptedModel, encryptedAesKey, iv } = encryptModel(modelFile, publicKeyBase64);
const { encryptedModel, encryptedAesKey, iv } = encryptModel(modelFile, publicKeyBase64);

const modelHash = await generateModelHash(modelFile); // Add await here
console.log("Model hash:", modelHash);
const signedHash = signModelHash(modelHash);
console.log("Signed hash:", signedHash);
res.status(200).json({
message: "Model encrypted and signed successfully",
encryptedModel: encryptedModel.toString("base64"),
encryptedAesKey: encryptedAesKey.toString("base64"),
iv: iv.toString("base64"),
signedHash: signedHash,
});
const modelHash = await generateModelHash(modelFile); // Add await here
console.log("Backend Model hash:", modelHash);
const signedHash = signModelHash(modelHash);
console.log("Bakend Signed Model hash:", signedHash);
res.status(200).json({
message: "Model encrypted and signed successfully",
encryptedModel: encryptedModel.toString("base64"),
encryptedAesKey: encryptedAesKey.toString("base64"),
iv: iv.toString("base64"),
signedHash: signedHash,
});
} catch (error) {
console.error("Error fetching and decrypting models:", error);
res.status(500).json({ error: "Failed to fetch and decrypt models." });
}
};

exports.getPublicVerificationKey = async (req, res, next) => {
try {
const publicKey = fs.readFileSync(path.join(__dirname, "../digital_signature_keys/public_key.pem"), "utf8");
res.status(200).json({ publicKey: publicKey });
} catch (error) {
console.error("Error fetching public key:", error);
res.status(500).json({ error: "Failed to fetch public key." });
}
}
try {
const publicKey = fs.readFileSync(path.join(__dirname, "../digital_signature_keys/public_key.pem"), "utf8");
res.status(200).json({ publicKey: publicKey });
} catch (error) {
console.error("Error fetching public key:", error);
res.status(500).json({ error: "Failed to fetch public key." });
}
}
2 changes: 1 addition & 1 deletion mern-backend/routes/api.js
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
const express = require("express");
const { getAllEncryptedModels,getPublicVerificationKey } = require("../controllers/modelController.js");
const { getAllEncryptedModels, getPublicVerificationKey } = require("../controllers/modelController.js");

const router = express.Router();

Expand Down
81 changes: 45 additions & 36 deletions mern-backend/utils/encryptionUtils.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,42 +3,51 @@ const fs = require("fs");
const path = require("path");

function encryptModel(modelFile, publicKeyBase64) {
if (!publicKeyBase64) {
throw new Error("Public key is required");
try {
if (!publicKeyBase64) {
throw new Error("Public key is required");
}

// Convert base64 to DER format
const publicKeyDer = Buffer.from(publicKeyBase64, 'base64');

// Create public key object using spki format
const publicKey = crypto.createPublicKey({
key: publicKeyDer,
format: 'der',
type: 'spki'
});

// Generate AES key and IV
const aesKey = crypto.randomBytes(32);
const iv = crypto.randomBytes(16);

// Encrypt the model with AES
const cipher = crypto.createCipheriv("aes-256-cbc", aesKey, iv);
let encryptedModel = Buffer.concat([
cipher.update(modelFile),
cipher.final()
]);

// Encrypt the AES key with RSA
const encryptedAesKey = crypto.publicEncrypt(
{
key: publicKey,
padding: crypto.constants.RSA_PKCS1_OAEP_PADDING,
oaepHash: 'sha256'
},
aesKey
);

return {
encryptedModel: encryptedModel.toString('base64'),
encryptedAesKey: encryptedAesKey.toString('base64'),
iv: iv.toString('base64')
};
} catch (error) {
console.error("Encryption error:", error);
throw error;
}

const publicKeyDer = Buffer.from(publicKeyBase64, 'base64');
const publicKey = crypto.createPublicKey({
key: publicKeyDer,
format: 'der',
type: 'spki'
});

console.log("Converted Public Key in PEM format:", publicKey);

const aesKey = crypto.randomBytes(32);

const iv = crypto.randomBytes(16);
const cipher = crypto.createCipheriv("aes-256-cbc", aesKey, iv);
let encryptedModel = cipher.update(modelFile);
encryptedModel = Buffer.concat([encryptedModel, cipher.final()]);

const encryptedAesKey = crypto.publicEncrypt(
{
key: publicKey,
padding: crypto.constants.RSA_PKCS1_OAEP_PADDING,
oaepHash: 'sha256'
},
aesKey
);

return {
encryptedModel,
encryptedAesKey,
iv
};
}

module.exports = {
encryptModel
};
module.exports = { encryptModel };
Loading

0 comments on commit f5b97e8

Please sign in to comment.