thesis-TA-resistant-anonymity-set/simulation/client.py

# -*- coding: utf-8 -*-

import json
import os
import sys

import requests

from simulation.encrypt_message import encrypt_message
from clustering.util import Constants


class Client(object):
    def __init__(self, server, text_file):
        self.server = server
        if not os.path.exists(text_file):
            raise IOError
        self.text_file = text_file

    def batch_upload_user_info(self, batch_size=100, gen_cover=0):

        with open(self.text_file, "r") as f:
            lines = f.readlines()
            batches = [
                lines[x : x + batch_size] for x in range(0, len(lines), batch_size)
            ]
            for batch in batches:
                    body = json.dumps(batch)
                    try:
                        body = encrypt_message(body)
                        r = requests.post(self.server, data=body)
                        print(r.text)
                        if gen_cover:
                            for _ in Constants.NUM_COVER:
                                cover_msg = self.cover_message_gen(body)
                                cover_msg = encrypt_message(cover_msg)
                                r = requests.post(self.server, data=cover_msg)
                                print(r.text)
                    except Exception as e:
                        print(e)

    def cover_message_gen(self, msg):
        user = msg.user_id
        time = msg.timestamp
        return {"timestamp": time, "user_id": user}


if __name__ == "__main__":

    server = sys.argv[1]
    text_file = sys.argv[2]
    client = Client(server, text_file)

    client.batch_upload_user_info(batch_size=10000)