refactored the whole project and revamped the ML model

Author: abheet19

.gitignore

@@ -0,0 +1,42 @@
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Virtual environments
+venv/
+env/
+ENV/
+env.bak/
+venv.bak/
+
+# Flask
+instance/
+.webassets-cache
+
+# Distribution / packaging
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# IDEs and editors
+.vscode/
+.idea/
+*.sublime-project
+*.sublime-workspace
+
+# OS files
+.DS_Store
+Thumbs.db

README.md

@@ -22,32 +22,72 @@ serializeDataCompany()
 ``` 
 ![Architecture](./architecture.png)
 
-Use the package manager [pip](https://pip.pypa.io/en/stable/) to install the prerequisites.
+### How to Run Locally
 
-```bash
-pip install -r requirements.txt
-```
+1. **Clone the Repository**
+   
+   ```bash
+   git clone https://github.com/yourusername/Partial-Homomorphic-Encryption-on-Machine-learning.git
+   cd Partial-Homomorphic-Encryption-on-Machine-learning
+   ```
+
+2. **Set Up a Conda Environment**
+   
+   ```bash
+   conda create -n phe_ml_env python=3.8
+   conda activate phe_ml_env
+   ```
+
+3. **Install Dependencies**
+   
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+4. **Train the Machine Learning Model**
+   
+   Before running the application, train the model to generate the `improved_model.pkl` file.
+   
+   ```bash
+   python train.py
+   ```
+
+5. **Generate Encryption Keys**
+   
+   ```python
+   # ...existing code...
+   python cust.py
+   ```
+
+6. **Start the Flask Application**
+   
+   ```bash
+   python app.py
+   ```
+
+7. **Access the Application**
+   
+   Open your web browser and navigate to `http://localhost:5000/` to interact with the application.
 
 ## Usage
 
-The main app.py (flask file) contains the code for implementing all the above functions
+The main `app.py` (Flask file) contains the code for implementing all the above functions
 ```bash
-@app.route('/customerEncryption',methods=['GET','POST']) 
-def customerEncryption():    #handles the encryption on client side
-
-@app.route('/company',methods=['GET','POST'])
-def company():               #handles the encryption on client side
+@app.route('/customerEncryption', methods=['GET','POST']) 
+def customerEncryption():    # Handles the encryption on client side
 
-@app.route('/result',methods=['GET','POST'])
-def result():                #generates the final output on the client side
+@app.route('/company', methods=['GET','POST'])
+def company():               # Handles the encryption on company side
 
+@app.route('/result', methods=['GET','POST'])
+def result():                # Generates the final output on the client side
 ```
 
-## Final thoughts
-* This is a high level glimpse to the future of ML and cryptography but not yet production ready.
-* Partial Homomorphic Encryption is a great way to perform complex ML computations on  encrypted sample
-* Evolving space with Fully-Homomorphic encryption which is in research  right now
-* This will likely be demand in the next 5-7 years.
+## Final Thoughts
+* This is a high-level glimpse into the future of ML and cryptography but not yet production-ready.
+* Partial Homomorphic Encryption is a great way to perform complex ML computations on encrypted samples.
+* Evolving space with Fully-Homomorphic Encryption which is in research right now.
+* This will likely be in demand in the next 5-7 years.
 
 ## Contributing
 Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.

app.py

@@ -1,56 +1,116 @@
+"""Flask application for Homomorphic Encryption-based Salary Prediction."""
+
 import numpy as np
-from flask import Flask, request, jsonify, render_template,redirect,url_for
+from flask import Flask, request, jsonify, render_template, redirect, url_for
 import json
+import os
 from cust import *
 from servercalc import *
 from os import path
 
 app = Flask(__name__)
 
-@app.route('/',methods=['GET','POST'])
-
+@app.route('/', methods=['GET', 'POST'])
 def home():
+    """
+    Render the home page.
     
+    Returns:
+        Rendered index.html template.
+    """
     return render_template('index.html')
 
-@app.route('/customerEncryption',methods=['GET','POST'])
+@app.route('/customerEncryption', methods=['GET', 'POST'])
 def customerEncryption():
+    """
+    Handle customer encryption data submission.
+    
+    - Generates encryption keys if not present.
+    - Encrypts customer data and saves it to data.json.
+    - Renders the cust.html template with encryption keys.
+    
+    Returns:
+        Rendered cust.html template with keys.
+    """
+    # Check if encryption keys exist; if not, generate and store them
     if not path.exists('custkeys.json'):
         storeKeys()
 
+    # Retrieve existing encryption keys
     pub_key, priv_key = getKeys()
 
+    # Extract features from the form and convert them to integers
     features = [int(x) for x in request.form.values()]
-    datafileCustomer=serializeDataCustomer(pub_key,features)
+    # Serialize and encrypt the customer data
+    datafileCustomer = serializeDataCustomer(pub_key, features)
+    # Save encrypted data to data.json
     with open('data.json', 'w') as file: 
         json.dump(datafileCustomer, file)
 
-    keys=[pub_key,priv_key]
-    return render_template('cust.html',keys=keys)
+    # Prepare keys to send to the frontend
+    keys = {'public_key': pub_key, 'private_key': priv_key}
+    # Render the customer details page with keys
+    return render_template('cust.html', keys=keys)
 
-@app.route('/company',methods=['GET','POST'])
+@app.route('/company', methods=['GET', 'POST'])
 def company():
-    datafileCompany=serializeDataCompany()
+    """
+    Handle company data submission for salary prediction.
+    
+    - Serializes and computes encrypted data.
+    - Saves the result to answer.json.
+    - Renders the company.html template with encrypted data.
+    
+    Returns:
+        Rendered company.html template with encrypted data.
+    """
+    # Serialize data for the company and compute encrypted results
+    datafileCompany = serializeDataCompany()
+    # Save encrypted company data to answer.json
     with open('answer.json', 'w') as file:
-        json.dump(datafileCompany,file)
-		
+        json.dump(datafileCompany, file)
 
-    return render_template('company.html',datafileCompany=datafileCompany)
+    # Render the company details page with encrypted data
+    return render_template('company.html', datafileCompany=datafileCompany)
 
-    #return redirect(url_for('result'))
+    # Optionally redirect to the result page
+    # return redirect(url_for('result'))
 
-@app.route('/result',methods=['GET','POST'])
+@app.route('/result', methods=['GET', 'POST'])
 def result():
-    answer_file=loadAnswer()
-    answer_key=paillier.PaillierPublicKey(n=int(answer_file['pubkey']['n']))
-    answer = paillier.EncryptedNumber(answer_key, int(answer_file['values'][0]), int(answer_file['values'][1]))
+    """
+    Process the encrypted data to decrypt and display the predicted salary.
+    
+    Steps:
+    - Load encrypted answer from answer.json.
+    - Decrypt the result using customer's private key.
+    - Round the final result to two decimal places.
+    - Render the result.html template with the final result.
+    
+    Returns:
+        Rendered result.html template with the final predicted salary.
+    """
+    # Load the encrypted answer data
+    answer_file = loadAnswer()
+    # Initialize the public key from the answer file
+    answer_key = paillier.PaillierPublicKey(n=int(answer_file['pubkey']['n']))
+    # Unpack the ciphertext and exponent
+    ciphertext_str, exponent = answer_file['values']  # Changed from answer_file['values'][0]
+    # Create an EncryptedNumber instance
+    answer = paillier.EncryptedNumber(answer_key, int(ciphertext_str), exponent)
+    # Retrieve customer's keys
     pub_key, priv_key = getKeys()
-    if (answer_key==pub_key):
-        final_result=priv_key.decrypt(answer)
- 	
+    if answer_key.n == pub_key.n:
+        # Decrypt the answer using the private key
+        final_result = priv_key.decrypt(answer)
+        # Round the result to two decimal places
+        final_result = round(final_result, 2)  # Rounded to 2 decimals
+    else:
+        final_result = "Keys do not match."
 
-    return render_template('result.html',final_result=final_result)
+    # Render the result page with the final predicted salary
+    return render_template('result.html', final_result=final_result)
 
-    
 if __name__ == "__main__":
-    app.run(debug=True)
+    # Run the Flask app with debug mode disabled for production
+    app.run(debug=False)

cust.py

@@ -1,46 +1,83 @@
+"""Customer-side encryption handling for Homomorphic Encryption-based Salary Prediction."""
+
 import phe as paillier
 import json
 
 def storeKeys():
-	public_key, private_key = paillier.generate_paillier_keypair()
-	keys={}
-	keys['public_key'] = {'n': public_key.n}
-	keys['private_key'] = {'p': private_key.p,'q':private_key.q}
-	with open('custkeys.json', 'w') as file: 
-		json.dump(keys, file)
+    """
+    Generate and store the Paillier public and private keys.
+    
+    The keys are saved to 'custkeys.json' in JSON format.
+    """
+    # Generate Paillier key pair
+    public_key, private_key = paillier.generate_paillier_keypair()
+    keys = {}
+    # Serialize public key
+    keys['public_key'] = {'n': public_key.n}
+    # Serialize private key components
+    keys['private_key'] = {'p': private_key.p, 'q': private_key.q}
+    # Save keys to custkeys.json
+    with open('custkeys.json', 'w') as file: 
+        json.dump(keys, file)
 
 def getKeys():
- 	with open('custkeys.json', 'r') as file:
-         
-		 keys=json.load(file)
-		 pub_key=paillier.PaillierPublicKey(n=int(keys['public_key']['n']))
-		 priv_key=paillier.PaillierPrivateKey(pub_key,keys['private_key']['p'],keys['private_key']['q'])
-		 return pub_key, priv_key 
+    """
+    Retrieve the Paillier public and private keys from 'custkeys.json'.
+    
+    Returns:
+        tuple: (PaillierPublicKey, PaillierPrivateKey)
+    """
+    with open('custkeys.json', 'r') as file:
+        keys = json.load(file)
+    # Reconstruct public key
+    pub_key = paillier.PaillierPublicKey(n=int(keys['public_key']['n']))
+    # Reconstruct private key
+    priv_key = paillier.PaillierPrivateKey(pub_key, int(keys['private_key']['p']), int(keys['private_key']['q']))
+    return pub_key, priv_key 
 
 def serializeDataCustomer(public_key, data):
- 	encrypted_data_list = [public_key.encrypt(x) for x in data]
- 	encrypted_data={}
- 	encrypted_data['public_key'] = {'n': public_key.n}
- 	encrypted_data['values'] = [(str(x.ciphertext()), x.exponent) for x in encrypted_data_list]
- 	serialized = json.dumps(encrypted_data)
- 	return serialized
+    """
+    Encrypt customer data using the provided public key.
+    
+    Args:
+        public_key (PaillierPublicKey): The public key for encryption.
+        data (list): List of integer features to encrypt.
+    
+    Returns:
+        dict: Serialized encrypted data containing public key and encrypted values.
+    """
+    # Encrypt each data feature
+    encrypted_data_list = [public_key.encrypt(x) for x in data]
+    encrypted_data = {}
+    # Store public key
+    encrypted_data['public_key'] = {'n': public_key.n}
+    # Store encrypted features as ciphertext and exponent tuples
+    encrypted_data['values'] = [(str(x.ciphertext()), x.exponent) for x in encrypted_data_list]
+    # Return as a dictionary instead of a JSON string
+    return encrypted_data
 
 def loadAnswer():
+    """
+    Load the encrypted answer data from 'answer.json'.
+    
+    Returns:
+        dict: Encrypted answer data.
+    """
     with open('answer.json', 'r') as file: 
-        ans=json.load(file)
-        answer=json.loads(ans)
-        return answer
+        ans = json.load(file)
+        return ans
 
+# Example usage (commented out)
 # pub_key, priv_key = getKeys()
 # data = age, he, al, gen = [24,4,6,1]      
 # serializeDataCustomer(pub_key, data)
-# datafileCustomer=serializeDataCustomer(pub_key, data)
+# datafileCustomer = serializeDataCustomer(pub_key, data)
 # with open('data.json', 'w') as file: 
 #     json.dump(datafileCustomer, file)
 
-# answer_file=loadAnswer()
-# answer_key=paillier.PaillierPublicKey(n=int(answer_file['pubkey']['n']))
+# answer_file = loadAnswer()
+# answer_key = paillier.PaillierPublicKey(n=int(answer_file['pubkey']['n']))
 # answer = paillier.EncryptedNumber(answer_key, int(answer_file['values'][0]), int(answer_file['values'][1]))
-# if (answer_key==pub_key):
+# if (answer_key == pub_key):
 #     print(priv_key.decrypt(answer))
- 	
+