Padra/oracle/worker/connector.py

import socket
import threading
import os
import config as cfg
from connection.socket import SocketServer
from connection.socket import SocketHTTP
from tls import models
import time
from tls import utils
import tls as t
from types import SimpleNamespace
import subprocess
from Crypto.Util.number import long_to_bytes, bytes_to_long
from tls.ghash import GhashCon
from Utils import *


class Connector:
    def __init__(self, parser, socket):
        self.parser = parser
        self.server = None
        self.serverDomain = None
        self.serverIP = None
        self.h = False
        self.s = SocketServer(socket, self.parser, "Connector")
        self.buffer = []
        self.handshaked = False
        gBackPrint("CONNECTOR", "")

        self.dataLen = 0
        self.auth_key = None
        self.last_key = None
        self.GhashCon = None
        self.key = None
        self.cipher = b""
        stream = os.popen("pwd")
        output = stream.read()
        self.path = output[:-1]
        self.t = threading.Thread(target=self.toCryptor)
        self.done = threading.Thread(target=self.done)
        self.doneTime = 5
        self.t.start()

    def done(self):
        try:
            while(self.doneTime > 0):
                time.sleep(1)
                self.doneTime = self.doneTime - 1
            self.server.close()
        except Exception as e:
            rBackPrint("An exception occurred done", e)

    def recv(self, data):
        try:
            self.buffer.append(data)
            if(self.h):
                self.doneTime = self.doneTime + 1
        except Exception as e:
            rBackPrint("An exception occurred recv", e)

    def oracleInitial(self, data):
        try:
            while len(self.s.connected) < 2:
                time.sleep(0.5)
            if len(self.s.connected) == 2:
                self.bearer = self.s.connected[1]
                self.cryptor = self.s.connected[0]
            self.key = data["key"]
            dataLen = data["dataLen"]
            if 0 == dataLen % 16:
                datalen = dataLen // 16
            else:
                datalen = (dataLen + (16 - dataLen % 16)) // 16

            self.dataLenBase = data["dataLen"]
            self.dataLen = datalen
            self.serverDomain = data["server"]
            self.serverIP = socket.gethostbyname(data["server"])
            bPrint("server IP", self.serverIP)
            if type(self.serverDomain) == str:
                self.serverDomain = [self.serverDomain]
            self.server = SocketHTTP(self.serverIP, 443, self)
        except Exception as e:
            rBackPrint("An exception occurred oracleInitial", e)

    def keyExchange(self, data):
        try:
            while not self.server:
                time.sleep(0.1)
            self.key_exchange = bytes.fromhex(data["keyExchange"])
            self.server.send(bytes.fromhex(self.client_hello()))
        except Exception as e:
            rBackPrint("An exception occurred keyExchange", e)

    def client_hello(self):
        try:
            print("Done 1")
            extensions = [
                models.ClientExtension.server_names(self.serverDomain),
                models.ClientExtension.supported_versions([b"\x03\x04"]),
                models.ClientExtension.key_share(
                    [models.KeyShareEntry(
                        models.NamedGroup.x25519, self.key_exchange)]
                ),
                models.ClientExtension.signature_algorithms(
                    list(models.SignatureScheme)
                ),
                models.ClientExtension.supported_groups(
                    list(models.NamedGroup)),
                models.ClientExtension.psk_key_exchange_modes(
                    list(models.PskKeyExchangeMode)
                ),
            ]

            compatibility_mode = True
            handshake = models.Handshake(
                models.HandshakeType.client_hello,
                models.ClientHello(
                    ...,
                    os.urandom(32),
                    os.urandom(32) if compatibility_mode else b"",
                    list(models.CipherSuite),
                    ...,
                    extensions,
                ),
            )
            self.client_hello_data = handshake.binary
            self.handshake_context = bytearray(self.client_hello_data)
            self.parser.parserSend(
                bytes(self.client_hello_data).hex(),
                self.cryptor,
                "fromContoCryHelloCLient",
            )

            data = models.ContentType.handshake.tls_plaintext(
                self.client_hello_data)
            return bytes.hex(data)
        except Exception as e:
            rBackPrint("An exception occurred client_hello", e)

    def toCryptor(self):
        try:
            while not self.h:
                while self.buffer and self.buffer[0] and len(self.buffer[0]) >= 10:
                    s = self.buffer[0][6:10]
                    l = int(s, 16) * 2 + 10
                    a = self.lenCalc(l)
                    if not self.handshaked:
                        self.parser.parserSend(a, self.cryptor, "fromContoCry")
        except Exception as e:
            rBackPrint("An exception occurred toCryptor", e)

    def lenCalc(self, len2):
        if self.buffer:
            if len(self.buffer[0]) > len2:
                a = self.buffer[0][:len2]
                self.buffer[0] = self.buffer[0][len2:]
                return a
            elif len(self.buffer[0]) == len2:
                return self.buffer.pop(0)
            elif len(self.buffer[0]) < len2:
                if self.buffer[1]:
                    self.buffer[0] = self.buffer[0] + self.buffer[1]
                    self.buffer.pop(1)
                    return self.lenCalc(len2)
                else:
                    self.lenCalc(len2)
        else:
            self.lenCalc(len2)

    def MPC(self):
        try:
            for i in range(self.dataLen * 4):
                print("send req to others")
                self.parser.parserSend(str(i), self.cryptor, "MPC")
                self.parser.parserSend(str(i), self.bearer, "MPC")
                s = [
                    self.path + "/MP-SPDZ/mascot-party.x",
                    "0",
                    "xor",
                    "-pn",
                    "12023",
                    "-h",
                    "192.168.178.105",
                    "-N",
                    "3",
                ]
                proc = subprocess.Popen(s, stdout=subprocess.PIPE)
                d = proc.stdout.readline()

                self.cipher = self.cipher + long_to_bytes(d, 4)
            self.cipher = self.cipher[: self.dataLenBase]
            head = b"\x17\x03\x03" + (self.dataLenBase + 16).to_bytes(2, "big")
            tag = self.GhashCon.ghash(head, self.cipher)
            tag = long_to_bytes(tag)
            bPrint("Tag", tag)
            Message = head + self.cipher + tag
            return Message
        except Exception as e:
            rBackPrint("An exception occurred MPC", e)

    def fromCrytoCon(self, data):
        try:
            self.server.send(bytes.fromhex(data["data"]))
        except Exception as e:
            rBackPrint("An exception occurred fromCrytoCon", e)

    def fromCrytoConDone(self, data):
        try:
            time.sleep(2)
            self.auth_key = data["auth_key"]
            self.last_key = data["last_key"]
            self.change = data["change"]
            self.finish = data["finish"]
            self.h = True

            gPrint("auth_key", self.auth_key)
            gPrint("last_key ", self.last_key)
            self.GhashCon = GhashCon(self.auth_key, self.last_key)
            m = self.MPC()
            self.done.start()
            self.server.send(bytes.fromhex(data["change"]))
            self.server.send(bytes.fromhex(data["finish"]))
            time.sleep(1)
            self.server.send((m))

        except Exception as e:
            rBackPrint("An exception occurred fromCrytoConDone", e)