#25 fix bugs + tests + refactoring

Kasean · Jul 5, 2024 · cb8be41 · cb8be41
1 parent 563e803
commit cb8be41
Show file tree

Hide file tree

Showing 14 changed files with 161 additions and 43 deletions.
diff --git a/artifact-processor/build.gradle b/artifact-processor/build.gradle
@@ -16,6 +16,9 @@ dependencies {
     implementation 'org.yaml:snakeyaml:2.0'
 
     implementation 'org.apache.kafka:kafka-clients:3.4.0'
+    implementation 'org.bouncycastle:bcprov-jdk18on:1.78.1'
+    implementation 'org.bouncycastle:bcpkix-jdk18on:1.78.1'
+
     testImplementation platform('org.junit:junit-bom:5.9.1')
     testImplementation 'org.junit.jupiter:junit-jupiter'
 }

diff --git a/artifact-processor/src/main/java/org/student/api/consumers/CreateConsumer.java b/artifact-processor/src/main/java/org/student/api/consumers/CreateConsumer.java
@@ -7,7 +7,7 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
-import org.student.api.ProcessMessageFunction;
+import org.student.api.utils.ProcessMessageFunction;
 
 import java.time.Duration;
 import java.util.Collections;

diff --git a/artifact-processor/src/main/java/org/student/api/consumers/DeleteConsumer.java b/artifact-processor/src/main/java/org/student/api/consumers/DeleteConsumer.java
@@ -6,7 +6,7 @@
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.StringDeserializer;
-import org.student.api.ProcessMessageFunction;
+import org.student.api.utils.ProcessMessageFunction;
 
 import java.time.Duration;
 import java.util.Collections;

diff --git a/artifact-processor/src/main/java/org/student/api/consumers/ReadConsumer.java b/artifact-processor/src/main/java/org/student/api/consumers/ReadConsumer.java
@@ -6,7 +6,7 @@
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.StringDeserializer;
-import org.student.api.ProcessMessageFunction;
+import org.student.api.utils.ProcessMessageFunction;
 
 import java.time.Duration;
 import java.util.Collections;

diff --git a/artifact-processor/src/main/java/org/student/api/factories/ConsumerFactory.java b/artifact-processor/src/main/java/org/student/api/factories/ConsumerFactory.java
@@ -1,6 +1,6 @@
 package org.student.api.factories;
 
-import org.student.api.OperationType;
+import org.student.api.utils.OperationType;
 import org.student.api.consumers.CreateConsumer;
 import org.student.api.consumers.DeleteConsumer;
 import org.student.api.consumers.MessageConsumer;

diff --git a/artifact-processor/src/main/java/org/student/api/managers/KafkaSendFacadeImpl.java b/artifact-processor/src/main/java/org/student/api/managers/KafkaSendFacadeImpl.java
@@ -12,7 +12,6 @@ public class KafkaSendFacadeImpl implements KafkaSendFacade {
     private final Map<String, MessageProducer> producers = new ConcurrentHashMap<>();
     private final String bootstrapServer;
 
-
     public KafkaSendFacadeImpl(String bootstrapServer) {
         this.bootstrapServer = bootstrapServer;
     }

diff --git a/...n/java/org/student/api/OperationType.java → .../org/student/api/utils/OperationType.java b/...n/java/org/student/api/OperationType.java → .../org/student/api/utils/OperationType.java
@@ -1,4 +1,4 @@
-package org.student.api;
+package org.student.api.utils;
 
 public enum OperationType {
     CREATE, READ, UPDATE, DELETE

diff --git a/...g/student/api/ProcessMessageFunction.java → ...ent/api/utils/ProcessMessageFunction.java b/...g/student/api/ProcessMessageFunction.java → ...ent/api/utils/ProcessMessageFunction.java
@@ -1,4 +1,4 @@
-package org.student.api;
+package org.student.api.utils;
 
 @FunctionalInterface
 public interface ProcessMessageFunction<F, S, T>{

diff --git a/artifact-processor/src/main/java/org/student/archiver/ArchiverService.java b/artifact-processor/src/main/java/org/student/archiver/ArchiverService.java
@@ -6,5 +6,5 @@ public interface ArchiverService {
 
     Artifact encrypt(byte[] rawArtifactMessage);
     Artifact decrypt(Artifact encryptedArtifact);
-
+    void shutdown();
 }
diff --git a/artifact-processor/src/main/java/org/student/archiver/ArchiverServiceImpl.java b/artifact-processor/src/main/java/org/student/archiver/ArchiverServiceImpl.java
@@ -1,5 +1,11 @@
 package org.student.archiver;
 
+import org.bouncycastle.asn1.x500.X500Name;
+import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
+import org.bouncycastle.cert.X509CertificateHolder;
+import org.bouncycastle.cert.X509v3CertificateBuilder;
+import org.bouncycastle.operator.ContentSigner;
+import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
 import org.student.configs.KeyStoreConfig;
 import org.student.models.Artifact;
 import org.student.models.ArtifactMetaInfo;
@@ -8,14 +14,19 @@
 import javax.crypto.Cipher;
 import javax.crypto.IllegalBlockSizeException;
 import javax.crypto.NoSuchPaddingException;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.OutputStream;
+import java.math.BigInteger;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.security.*;
 import java.security.cert.Certificate;
 import java.security.cert.CertificateException;
+import java.security.cert.CertificateFactory;
+import java.util.Date;
 import java.util.Optional;
 import java.util.UUID;
 
@@ -27,6 +38,7 @@ public class ArchiverServiceImpl implements ArchiverService {
     private static final String RSA = "RSA";
     private static final String ALIAS_PATTERN = "%s%s";
     private static final String SKIP_DECRYPTION = " Skip decryption.";
+    private static final String KEYSTORE_JKS_ROOT_PATH = "/keystore.jks";
 
     private final KeyStore keyStore;
     private final char[] rootPassword;
@@ -60,11 +72,9 @@ public Artifact encrypt(byte[] rawArtifactMessage) {
 
         PublicKey publicKey = keyPair.getPublic();
 
-        Cipher cipher = createCipher();
+        Cipher cipher = createCipher(Cipher.ENCRYPT_MODE, publicKey);
         if (cipher == null) return artifact;
 
-        if (!initCipher(cipher, publicKey)) return artifact;
-
         return encryptByteArray(cipher, artifact);
     }
 
@@ -90,7 +100,19 @@ private boolean loadKeyStore(KeyStore keyStore) {
     }
 
     private boolean addPrivateKeyToKeyStore(KeyStore keyStore, KeyPair keyPair, ArtifactMetaInfo metaInfo) {
-        KeyStore.PrivateKeyEntry privateKeyEntry = new KeyStore.PrivateKeyEntry(keyPair.getPrivate(), new Certificate[]{});
+
+        long validity = 100 * 365 * 24 * 60 * 60; // 100 years
+        String sigAlgName = "SHA256WithRSA";
+
+        Certificate certificate = null;
+        try {
+            certificate = generateSelfSignedCertificate(keyPair, "CN=Test", validity, sigAlgName);
+        } catch (Exception e) {
+            System.err.println("Error creating certificate: " + e.getMessage() + e + SKIP_ENCRYPTION);
+            return false;
+        }
+
+        KeyStore.PrivateKeyEntry privateKeyEntry = new KeyStore.PrivateKeyEntry(keyPair.getPrivate(), new Certificate[]{certificate});
         try {
             var alias = generateAlias();
             keyStore.setEntry(alias, privateKeyEntry, new KeyStore.PasswordProtection(rootPassword));
@@ -102,13 +124,34 @@ private boolean addPrivateKeyToKeyStore(KeyStore keyStore, KeyPair keyPair, Arti
         }
     }
 
+    private Certificate generateSelfSignedCertificate(KeyPair keyPair, String dn, long validity, String sigAlgName) throws Exception {
+        X500Name issuerName = new X500Name(dn);
+        BigInteger serial = BigInteger.valueOf(new SecureRandom().nextInt());
+
+        Date from = new Date();
+        Date to = new Date(from.getTime() + validity * 1000);
+
+        SubjectPublicKeyInfo publicKeyInfo = SubjectPublicKeyInfo.getInstance(keyPair.getPublic().getEncoded());
+
+        X509v3CertificateBuilder certBuilder = new X509v3CertificateBuilder(issuerName, serial, from, to, issuerName, publicKeyInfo);
+
+        ContentSigner signer = new JcaContentSignerBuilder(sigAlgName).build(keyPair.getPrivate());
+
+        X509CertificateHolder certHolder = certBuilder.build(signer);
+
+        CertificateFactory certFactory = CertificateFactory.getInstance("X.509");
+
+        return certFactory.generateCertificate(new ByteArrayInputStream(certHolder.getEncoded()));
+    }
+
     private String generateAlias() {
         return String.format(ALIAS_PATTERN, defaultRSAAlias, UUID.randomUUID());
     }
 
     private boolean saveKeyStoreToDisk(KeyStore keyStore) {
         try {
-            Path path = Paths.get(pathToKeyStore + "/keystore.jks");
+            Path path = Paths.get(pathToKeyStore + KEYSTORE_JKS_ROOT_PATH);
+            Files.createDirectories(path.getParent());
             try (OutputStream fos = Files.newOutputStream(path)) {
                 keyStore.store(fos, rootPassword);
             }
@@ -119,25 +162,17 @@ private boolean saveKeyStoreToDisk(KeyStore keyStore) {
         }
     }
 
-    private Cipher createCipher() {
+    private Cipher createCipher(int mode, Key key) {
         try {
-            return Cipher.getInstance(RSA);
-        } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
+            var instance =  Cipher.getInstance(RSA);
+            instance.init(mode, key);
+            return instance;
+        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException e) {
             System.err.println("Error in instancing Cipher:" + e.getMessage() + e + SKIP_ENCRYPTION);
             return null;
         }
     }
 
-    private boolean initCipher(Cipher cipher, PublicKey publicKey) {
-        try {
-            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
-            return true;
-        } catch (InvalidKeyException e) {
-            System.err.println("Error in initing Cipher:" + e.getMessage() + e + SKIP_ENCRYPTION);
-            return false;
-        }
-    }
-
     private Artifact encryptByteArray(Cipher cipher, Artifact artifact) {
         try {
             var finalRes = cipher.doFinal(artifact.getArtifactData());
@@ -154,15 +189,12 @@ public Artifact decrypt(Artifact encryptedArtifact) {
 
         if (keyStore == null) return encryptedArtifact;
 
-        if (!loadKeyStore(keyStore)) return encryptedArtifact;
         PrivateKey privateKey = getPrivateKeyFromKeyStore(keyStore, archiverRepository.getArtifactAlias(encryptedArtifact.getMetaInfo()));
         if (privateKey == null) return encryptedArtifact;
 
-        Cipher cipher = createCipher();
+        Cipher cipher = createCipher(Cipher.DECRYPT_MODE, privateKey);
         if (cipher == null) return encryptedArtifact;
 
-        if (!initCipherForDecryption(cipher, privateKey)) return encryptedArtifact;
-
         var decrypted = decryptByteArray(cipher, encryptedArtifact);
         if (decrypted.length == 0) return encryptedArtifact;
 
@@ -171,6 +203,11 @@ public Artifact decrypt(Artifact encryptedArtifact) {
         return encryptedArtifact;
     }
 
+    @Override
+    public void shutdown() {
+        saveKeyStoreToDisk(keyStore);
+    }
+
     private PrivateKey getPrivateKeyFromKeyStore(KeyStore keyStore, Optional<String> artifactAliasOpt) {
 
         if (artifactAliasOpt.isEmpty()){
@@ -181,24 +218,21 @@ private PrivateKey getPrivateKeyFromKeyStore(KeyStore keyStore, Optional<String>
         var artifactAlias = artifactAliasOpt.get();
 
         try {
-            KeyStore.PrivateKeyEntry privateKeyEntry = (KeyStore.PrivateKeyEntry) keyStore.getEntry(artifactAlias, new KeyStore.PasswordProtection(rootPassword));
-            return privateKeyEntry.getPrivateKey();
+
+            Key key = keyStore.getKey(artifactAlias, rootPassword);
+            if (key instanceof PrivateKey) {
+                PrivateKey privateKey = (PrivateKey) key;
+                return privateKey;
+            } else {
+                System.err.println("Error in getting private key from KeyStore: not founded key in store." + SKIP_ENCRYPTION);
+                return null;
+            }
         } catch (NoSuchAlgorithmException | UnrecoverableEntryException | KeyStoreException e) {
-            System.err.println("Error in getting private key from KeyStore:" + e.getMessage() + e + SKIP_DECRYPTION);
+            System.err.println("Error in getting private key from KeyStore: " + e.getMessage() + e + SKIP_DECRYPTION);
             return null;
         }
     }
 
-    private boolean initCipherForDecryption(Cipher cipher, PrivateKey privateKey) {
-        try {
-            cipher.init(Cipher.DECRYPT_MODE, privateKey);
-            return true;
-        } catch (InvalidKeyException e) {
-            System.err.println("Error in initing Cipher for decryption:" + e.getMessage() + e + SKIP_DECRYPTION);
-            return false;
-        }
-    }
-
     private byte[] decryptByteArray(Cipher cipher, Artifact encryptedMessage) {
         try {
             return cipher.doFinal(encryptedMessage.getArtifactData());

diff --git a/artifact-processor/src/main/java/org/student/archiver/KeyStoreManager.java b/artifact-processor/src/main/java/org/student/archiver/KeyStoreManager.java
@@ -35,12 +35,18 @@ private static void loadKSFromDisk(char[] rootPassword, String pathToKeyStore) t
         try {
             Path path = Paths.get(pathToKeyStore + "/keystore.jks");
             try (InputStream fis = Files.newInputStream(path)) {
+                instance = KeyStore.getInstance("JKS");
                 instance.load(fis, rootPassword);
             }
         } catch (IOException | NoSuchAlgorithmException | CertificateException e) {
             System.err.println("Error in loading KeyStore from ROM:" + e.getMessage() + e + " Create new instance.");
 
             instance = KeyStore.getInstance("JKS");
+            try {
+                instance.load(null, rootPassword);
+            } catch (IOException | CertificateException | NoSuchAlgorithmException ex) {
+                System.err.println("Key store initialization failed: " + e.getMessage() + e);
+            }
         }
     }
 

diff --git a/artifact-processor/src/main/java/org/student/services/ArtifactsService.java b/artifact-processor/src/main/java/org/student/services/ArtifactsService.java
@@ -11,4 +11,6 @@ public interface ArtifactsService {
     void updateArtifactMessage(String key, UUID id, byte[] newArtifactMessage); // TODO: not in alfa
     void deleteArtifactMessage(String key, UUID id, String topic);
 
+    void shutdown();
+
 }
diff --git a/artifact-processor/src/main/java/org/student/services/ArtifactsServiceImpl.java b/artifact-processor/src/main/java/org/student/services/ArtifactsServiceImpl.java
@@ -71,4 +71,9 @@ public void deleteArtifactMessage(String key, UUID id, String topic) {
 
         producer.send(topic, key, BodyArtifactMessage.class, artifactResponse);
     }
+
+    @Override
+    public void shutdown() {
+        archiver.shutdown();
+    }
 }
diff --git a/artifact-processor/src/test/java/archiver/ArchiverServiceTest.java b/artifact-processor/src/test/java/archiver/ArchiverServiceTest.java
@@ -0,0 +1,69 @@
+package archiver;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.student.archiver.ArchiverService;
+import org.student.archiver.ArchiverServiceImpl;
+import org.student.configs.KeyStoreConfig;
+import org.student.models.Artifact;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class ArchiverServiceTest {
+
+    private static final String RSA_ALIAS = "RSAAlias";
+    private static final String ROOT_PASSWORD = "P@55w0rd";
+    private static final String PATH_TO_KEY_STORE = "src/test/resources";
+    private ArchiverService archiverService;
+
+    @BeforeEach
+    void setUp() {
+
+        KeyStoreConfig config = new KeyStoreConfig();
+        config.setDefaultRSAAlias(RSA_ALIAS);
+        config.setRootPassword(ROOT_PASSWORD);
+        config.setPathToKeyStore(PATH_TO_KEY_STORE);
+
+        archiverService = new ArchiverServiceImpl(config);
+    }
+
+    @AfterEach
+    void tearDown() {
+        Path path = Paths.get(PATH_TO_KEY_STORE + "/keystore.jks");
+        try {
+            Files.deleteIfExists(path);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
+
+    @Test
+    void testEncrypt() {
+        byte[] rawArtifactMessage = "Test message".getBytes();
+        Artifact artifact = archiverService.encrypt(rawArtifactMessage);
+
+        assertNotNull(artifact, "Artifact should not be null");
+        assertNotEquals(rawArtifactMessage[0], artifact.getArtifactData()[0]);
+    }
+
+    @Test
+    void testDecrypt() {
+        var testStr = "Test message";
+        byte[] rawArtifactMessage = testStr.getBytes();
+        Artifact encryptedArtifact = archiverService.encrypt(rawArtifactMessage);
+        Artifact decryptedArtifact = archiverService.decrypt(encryptedArtifact);
+
+        var actualStr = new String(decryptedArtifact.getArtifactData());
+
+        assertNotNull(decryptedArtifact, "Decrypted artifact should not be null");
+        assertEquals(actualStr, testStr);
+
+    }
+
+}