Hybrid encryption for API operations

import requests as req
import json
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher
import base64
from datetime import datetime
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad
aes_key = get_random_bytes(16)
aes_iv = get_random_bytes(16)


max_chunk_size = 53
def encrypt(s, pk):
    key = RSA.import_key(base64.b64decode(pk))
    cipher = PKCS1_cipher.new(key)
    data_to_encrypt = bytes(s.encode('utf8'))
    chunks = [data_to_encrypt[i:i + max_chunk_size] for i in range(0, len(data_to_encrypt), max_chunk_size)]
    encrypted_chunks = [cipher.encrypt(chunk) for chunk in chunks]
    encrypted = b''.join(encrypted_chunks)
    encoded_encrypted_chunks = base64.b64encode(encrypted).decode('utf-8')
    return encoded_encrypted_chunks

def encrypt_body(body, pk):
    cipher = AES.new(aes_key, AES.MODE_CBC, aes_iv)
    return base64.b64encode(cipher.encrypt(pad(bytes(body.encode('utf8')), AES.block_size))).decode('utf8')

class Header:
    auth: str
    ts: str
    uri: str
    keys: bytes
    def __init__(self, auth, ts, uri, keys):
        self.auth = auth
        self.ts = ts
        self.uri = uri
        self.keys = keys

def serialize_struct(struct):
    return json.dumps({
        'auth': struct.auth,
        'ts': struct.ts,
        'uri': struct.uri,
        'keys': base64.b64encode(struct.keys).decode('utf-8')
    })

ip = '10.10.10.1' # The IP address of the OBServer node that sends the request.
port = 2886 # The port number of the OBServer node that sends the request.
pwd = '******' # The password for the root user of the sys tenant in the OceanBase cluster.
uri = '/api/v1/ob/init' # The URI of the request.

body = {} # Update the request body.


# Get the public key.
resp = req.get(f'http://{ip}:{port}/api/v1/secret').text
resp = json.loads(resp)
pk = resp['data']['public_key']
header = Header(pwd, str(int(datetime.now().timestamp()) + 100000), uri, aes_key + aes_iv)


print("request headers: ", 'X-OCS-Header: ' + encrypt(serialize_struct(header), pk))
print("\n")
print("request body: ", encrypt_body(json.dumps(body), pk))
print("\n")
print("encrypt password: ", encrypt(pwd, pk))

package main

import (
    "bytes"
    "crypto/aes"
    "crypto/cipher"
    "crypto/rand"
    "crypto/rsa"
    "crypto/x509"
    "encoding/base64"
    "encoding/json"
    "fmt"
    "io"
    "net/http"
    "time"
)

var (
    ip   = "10.10.10.1"  // The IP address of the OBServer node that sends the request.
    port = 2886   // The port number of the OBServer node that sends the request.
    uri  = "/api/v1/ob/init"   // The URI of the request.
    pwd  = "******"   // The password for the root user of the sys tenant in the OceanBase cluster.
    body = map[string]interface{}{} // Update the request body.
)
var pk string

func main() {
    var err error
    pk, err = getPublicKey(fmt.Sprintf("http://%s:%d/api/v1/secret", ip, port))
    if err != nil {
        fmt.Printf("get public key failed: %s", err.Error())
        return
    }
    fmt.Printf("target agent's public key: %s\n\n", pk)
    encryptedBody, header, err := BuildBodyAndHeader(uri, body)
    if err != nil {
        return
    }
    fmt.Printf("request header: %v\n\nrequest body: %v\n", header, encryptedBody)
}

// getPublicKey function retrieves the public key from the API
func getPublicKey(url string) (string, error) {
    resp, err := http.Get(url)
    if err != nil {
        return "", err
    }
    defer resp.Body.Close()

    body, err := io.ReadAll(resp.Body)
    if err != nil {
        return "", err
    }

    var response struct {
        Data struct {
            PublicKey string `json:"public_key"`
        } `json:"data"`
    }
    if err = json.Unmarshal(body, &response); err != nil {
        return "", err
    }
    return response.Data.PublicKey, nil
}

func BuildBodyAndHeader(uri string, param interface{}) (encryptedBody interface{}, header map[string]string, err error) {
    encryptedBody, Key, Iv, err := EncryptBodyWithAes(param)
    if err != nil {
        return nil, nil, err
    }

    header = BuildHeader(Key, Iv)
    return encryptedBody, header, nil
}

type HttpHeader struct {
    Auth  string
    Ts    string
    Uri   string
    Keys  []byte
}

func BuildHeader(keys ...[]byte) map[string]string {
    headers := make(map[string]string)

    var aesKeys []byte
    if len(keys) != 2 {
        aesKeys = nil
    } else {
        aesKeys = append(keys[0], keys[1]...)
    }
    header := HttpHeader{
        Auth: pwd,
        Ts:   fmt.Sprintf("%d", time.Now().Add(100*time.Second).Unix()),
        Uri:  uri,
        Keys: aesKeys,
    }
    mAuth, err := json.Marshal(header)
    if err != nil {
        return nil
    }
    encrypt, err := RSAEncrypt(mAuth, pk)
    if err != nil {
        return nil
    }
    headers["X-OCS-Header"] = encrypt
    return headers
}

func RSAEncrypt(raw []byte, pk string) (string, error) {
    pkix, err := base64.StdEncoding.DecodeString(pk)
    if err != nil {
        return "", err
    }
    pub, err := x509.ParsePKCS1PublicKey(pkix)
    if err != nil {
        return "", err
    }
    if len(raw) == 0 {
        b, err := rsa.EncryptPKCS1v15(rand.Reader, pub, raw)
        return base64.StdEncoding.EncodeToString(b), err
    }
    // Encrypt data segments.
    blockSize := 512/8 - 11
    numBlocks := (len(raw) + blockSize - 1) / blockSize
    ciphertext := make([]byte, 0)
    for i := 0; i < numBlocks; i++ {
        start := i * blockSize
        end := (i + 1) * blockSize
        if end > len(raw) {
            end = len(raw)
        }
        encrypted, err := rsa.EncryptPKCS1v15(rand.Reader, pub, raw[start:end])
        if err != nil {
            return "", err
        }
        ciphertext = append(ciphertext, encrypted...)
    }
    return base64.StdEncoding.EncodeToString(ciphertext), err
}

func EncryptBodyWithAes(body interface{}) (encryptedBody interface{}, key []byte, iv []byte, err error) {
    if body == nil {
        return
    }
    mBody, err := json.Marshal(body)
    if err != nil {
        return
    }
    key = make([]byte, 16)
    iv = make([]byte, 16) // equal to block_size, 16 bytes
    _, err = rand.Read(key)
    if err != nil {
        return
    }
    _, err = rand.Read(iv)
    if err != nil {
        return
    }
    encryptedBody, err = AESEncrypt(mBody, key, iv)
    return
}

func AESEncrypt(raw []byte, key []byte, iv []byte) (string, error) {
    // Create an AES encryptor.
    block, err := aes.NewCipher(key)
    if err != nil {
        return "", err
    }
    raw = pkcs5Padding(raw, block.BlockSize())
    // Create an AES encryptor in CBC mode.
    mode := cipher.NewCBCEncrypter(block, iv)

    // Encrypt data.
    ciphertext := make([]byte, len(raw))
    mode.CryptBlocks(ciphertext, raw)

    return base64.StdEncoding.EncodeToString(ciphertext), err
}

func pkcs5Padding(ciphertext []byte, blockSize int) []byte {
    padding := blockSize - len(ciphertext)%blockSize
    padtext := bytes.Repeat([]byte{byte(padding)}, padding)
    return append(ciphertext, padtext...)
}

package com.oceanbase;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.Serializable;
import java.math.BigInteger;
import java.net.HttpURLConnection;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.RSAPublicKeySpec;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.net.URL;


import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.ASN1Sequence;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;

import com.google.gson.Gson;

public class Main {
    /* Modify the password, URI, and target agent. */
    private static final String pwd = "******";  // The password for the root user of the sys tenant in the OceanBase cluster.
    private static final String uri = "/api/v1/obcluster/config"; // The URI of the request.
    private static final String agent = "10.10.10.1:2886";  // The IP address and port number of the OBServer node that sends the request.

    private static final String RSA_ALGORITHM = "RSA";
    private static final Integer RSA_KEY_SIZE = 512;
    private static final String RSA_TRANSFORMATION = "RSA/ECB/PKCS1Padding";
    private static final String AES_ALGORITHM = "AES";
    private static final Integer AES_KEY_SIZE = 128;
    private static final String AES_TRANSFORMATION = "AES/CBC/PKCS5Padding";

    public static class HTTPHeader implements Serializable {
        private static final long serialVersionUID = 1L;
        public String auth;
        public String ts;
        public String uri;
        private String keys;

        public HTTPHeader(String pwd, int timestamp, String uri, byte[] keys) {
            this.auth = pwd;
            this.ts = String.valueOf(timestamp);
            this.uri = uri;
            this.keys = Base64.getEncoder().encodeToString(keys);
        }

    }

    public static String serializeStruct(HTTPHeader header) {
        Gson gson = new Gson();

        // Convert the struct into the JSON format and perform Base64 encoding on the keys field.
        return gson.toJson(header);
    }

    public static String encryptRSAWithSegment(Cipher cipher, byte[] data) throws Exception {
        int inputLen = data.length;
        int maxSize = RSA_KEY_SIZE / 8 - 11; // The key length, which is equal to the data length minus the padding length.
        int offSet = 0;
        byte[] cache;
        int i = 0;
        List<byte[]> encryptedSegments = new ArrayList<>();

        // Encrypt data segments.
        while (inputLen - offSet > 0) {
            if (inputLen - offSet > maxSize) {
                cache = cipher.doFinal(data, offSet, maxSize);
            } else {
                cache = cipher.doFinal(data, offSet, inputLen - offSet);
            }
            encryptedSegments.add(cache);
            i++;
            offSet = i * maxSize;
        }

        // Combine encrypted segments.
        int totalLen = encryptedSegments.stream().mapToInt(segment -> segment.length).sum();
        byte[] encryptedData = new byte[totalLen];
        int currentIndex = 0;
        for (byte[] segment : encryptedSegments) {
            System.arraycopy(segment, 0, encryptedData, currentIndex, segment.length);
            currentIndex += segment.length;
        }

        return Base64.getEncoder().encodeToString(encryptedData);
    }

    private static String getPublicKey(String urlStr) throws Exception {
        URL url = new URL(urlStr);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod("GET");

        int responseCode = conn.getResponseCode();
        if (responseCode == HttpURLConnection.HTTP_OK) {
            try (BufferedReader in = new BufferedReader(new InputStreamReader(conn.getInputStream()))) {
                String inputLine;
                StringBuilder response = new StringBuilder();

                while ((inputLine = in.readLine()) != null) {
                    response.append(inputLine);
                }

                System.out.println("Response: " + response.toString());
                JSONObject jsonResponse = JSON.parseObject(response.toString());
                System.out.println("Public Key: " + jsonResponse.getJSONObject("data").getString("public_key"));
                return jsonResponse.getJSONObject("data").getString("public_key");
            } catch (Exception e){
                System.out.println("Failed to read response: " + e.getMessage());
                throw new Exception("Failed to read response", e);
            }
        } else {
            throw new Exception("HTTP request failed with code " + responseCode);
        }
    }

    private static PublicKey convertPKCS1ToPublicKey(String pkcs1PublicKeyStr) throws Exception {
        pkcs1PublicKeyStr = pkcs1PublicKeyStr.replaceAll("\\n", "")
                .replace("-----BEGIN RSA PUBLIC KEY-----", "")
                .replace("-----END RSA PUBLIC KEY-----", "");
        byte[] pkcs1PublicKey = Base64.getDecoder().decode(pkcs1PublicKeyStr);

        ASN1InputStream asn1InputStream = new ASN1InputStream(pkcs1PublicKey);
        ASN1Sequence sequence = (ASN1Sequence) asn1InputStream.readObject();
        BigInteger modulus = ((ASN1Integer) sequence.getObjectAt(0)).getValue();
        BigInteger publicExponent = ((ASN1Integer) sequence.getObjectAt(1)).getValue();
        asn1InputStream.close();

        RSAPublicKeySpec publicKeySpec = new RSAPublicKeySpec(modulus, publicExponent);

        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePublic(publicKeySpec);
    }


    public static void main(String[] args) throws Exception {
        // Get the public key.
        String publicKeyStr = getPublicKey("http://" +agent+ "/api/v1/secret");
        PublicKey publicKey = convertPKCS1ToPublicKey(publicKeyStr);

        // Initialize the RSA and AES keys.
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance(RSA_ALGORITHM);
        keyGen.initialize(RSA_KEY_SIZE); // The RSA key, which contains 2048 bits.

        KeyGenerator aesKeyGen = KeyGenerator.getInstance(AES_ALGORITHM);
        aesKeyGen.init(AES_KEY_SIZE); // The AES key, which contains 128 bits.
        SecretKey aesKey = aesKeyGen.generateKey();
        byte[] aesKeyBytes = aesKey.getEncoded();
        byte[] iv = new byte[16]; // The IV, which contains 16 bytes.
        SecureRandom random = new SecureRandom();
        random.nextBytes(iv);

        // Encrypt data segments by using the RSA algorithm.
        Cipher rsaCipher = Cipher.getInstance(RSA_TRANSFORMATION);
        rsaCipher.init(Cipher.ENCRYPT_MODE, publicKey);


        // Get the URI, timestamp, AES key, and IV to be added to the headers.
        // Perform Base64 encoding on the data and concatenate the data into a string.
        long timestamp = System.currentTimeMillis() / 1000 + 100000;
        HTTPHeader header = new HTTPHeader(pwd, (int) timestamp, uri, concatenate(aesKeyBytes, iv));
        System.out.println("Header: " +  serializeStruct(header));
        String encryptedHeader = encryptRSAWithSegment(rsaCipher, serializeStruct(header).getBytes(StandardCharsets.UTF_8));

        // Encrypt the HTTP struct by using the AES algorithm.
        Cipher aesCipher = Cipher.getInstance(AES_TRANSFORMATION);
        aesCipher.init(Cipher.ENCRYPT_MODE, aesKey, new IvParameterSpec(iv));

        /* Construct the request body. */
        JSONObject databody = new JSONObject();
        data.body.put("clusterId", 1);
        data.body.put("clusterName", "21421");
        data.body.put("rootPwd", "2145325");
        System.out.println(data.body.toString());
        byte[] encryptedBody = aesCipher.doFinal(data.body.toString().getBytes(StandardCharsets.UTF_8));
        String encryptedBodyBase64 = Base64.getEncoder().encodeToString(encryptedBody);

        // Return the result.
        System.out.println("Encrypted Header (Base64, RSA-Encrypted): " + encryptedHeader);
        System.out.println("Encrypted Body (Base64, AES-Encrypted): " + encryptedBodyBase64);
    }

    private static byte[] concatenate(byte[]... arrays) {
        int totalLength = 0;
        for (byte[] array : arrays) {
            totalLength += array.length;
        }
        byte[] result = new byte[totalLength];
        int currentIndex = 0;
        for (byte[] array : arrays) {
            System.arraycopy(array, 0, result, currentIndex, array.length);
            currentIndex += array.length;
        }
        return result;
    }
}