Skip to content

Commit

Permalink
NIFI-8251 Refactored Public Key ID Property to Public Key Search
Browse files Browse the repository at this point in the history
  • Loading branch information
exceptionfactory committed Feb 27, 2021
1 parent ae65672 commit 8040dee
Show file tree
Hide file tree
Showing 5 changed files with 48 additions and 46 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@
import org.apache.nifi.processors.pgp.attributes.FileEncoding;
import org.apache.nifi.processors.pgp.attributes.SymmetricKeyAlgorithm;
import org.apache.nifi.processors.pgp.exception.PGPEncryptionException;
import org.apache.nifi.processors.pgp.exception.PGPProcessException;
import org.apache.nifi.stream.io.StreamUtils;
import org.apache.nifi.util.StringUtils;

Expand Down Expand Up @@ -140,10 +139,10 @@ public class EncryptContentPGP extends AbstractProcessor {
.identifiesControllerService(PGPPublicKeyService.class)
.build();

public static final PropertyDescriptor PUBLIC_KEY_ID = new PropertyDescriptor.Builder()
.name("public-key-id")
.displayName("Public Key ID")
.description("PGP Public Key Identifier for encryption formatted as uppercase hexadecimal string of 16 characters")
public static final PropertyDescriptor PUBLIC_KEY_SEARCH = new PropertyDescriptor.Builder()
.name("public-key-search")
.displayName("Public Key Search")
.description("PGP Public Key Search will be used to match against the User ID or Key ID when formatted as uppercase hexadecimal string of 16 characters")
.expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES)
.addValidator(StandardValidators.ATTRIBUTE_EXPRESSION_LANGUAGE_VALIDATOR)
.dependsOn(PUBLIC_KEY_SERVICE)
Expand All @@ -154,8 +153,6 @@ public class EncryptContentPGP extends AbstractProcessor {

private static final int OUTPUT_BUFFER_SIZE = 8192;

private static final int HEXADECIMAL = 16;

private static final Set<Relationship> RELATIONSHIPS = new HashSet<>(Arrays.asList(SUCCESS, FAILURE));

private static final List<PropertyDescriptor> DESCRIPTORS = Arrays.asList(
Expand All @@ -164,7 +161,7 @@ public class EncryptContentPGP extends AbstractProcessor {
FILE_ENCODING,
PASSPHRASE,
PUBLIC_KEY_SERVICE,
PUBLIC_KEY_ID
PUBLIC_KEY_SEARCH
);

/**
Expand Down Expand Up @@ -271,34 +268,24 @@ private List<PGPKeyEncryptionMethodGenerator> getEncryptionMethodGenerators(fina
generators.add(new JcePBEKeyEncryptionMethodGenerator(passphrase).setSecureRandom(secureRandom));
}

final PropertyValue publicKeyIdProperty = context.getProperty(PUBLIC_KEY_ID);
if (publicKeyIdProperty.isSet()) {
final String publicKeyId = publicKeyIdProperty.evaluateAttributeExpressions(flowFile).getValue();
getLogger().debug("Requested Public Key ID [{}]", publicKeyId);
final PropertyValue publicKeySearchProperty = context.getProperty(PUBLIC_KEY_SEARCH);
if (publicKeySearchProperty.isSet()) {
final String publicKeySearch = publicKeySearchProperty.evaluateAttributeExpressions(flowFile).getValue();
getLogger().debug("Public Key Search [{}]", publicKeySearch);

final long keyId = parsePublicKeyId(publicKeyId);
final PGPPublicKeyService publicKeyService = context.getProperty(PUBLIC_KEY_SERVICE).asControllerService(PGPPublicKeyService.class);
final Optional<PGPPublicKey> optionalPublicKey = publicKeyService.findPublicKey(keyId);
final Optional<PGPPublicKey> optionalPublicKey = publicKeyService.findPublicKey(publicKeySearch);
if (optionalPublicKey.isPresent()) {
final PGPPublicKey publicKey = optionalPublicKey.get();
generators.add(new BcPublicKeyKeyEncryptionMethodGenerator(publicKey).setSecureRandom(secureRandom));
} else {
throw new PGPEncryptionException(String.format("Public Key ID [%s] not found", publicKeyId));
throw new PGPEncryptionException(String.format("Public Key not found using search [%s]", publicKeySearch));
}
}

return generators;
}

private long parsePublicKeyId(final String publicKeyId) {
try {
return Long.parseUnsignedLong(publicKeyId, HEXADECIMAL);
} catch (final RuntimeException e) {
final String message = String.format("Parsing Public Key ID [%s] Failed", publicKeyId);
throw new PGPProcessException(message, e);
}
}

private SymmetricKeyAlgorithm getSymmetricKeyAlgorithm(final ProcessContext context) {
final String algorithm = context.getProperty(SYMMETRIC_KEY_ALGORITHM).getValue();
return SymmetricKeyAlgorithm.valueOf(algorithm);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,8 @@ public void testSuccessPasswordBasedEncryptionFileEncodingAscii() throws IOExcep
public void testSuccessPublicKeyEncryptionRsaPublicKey() throws IOException, InitializationException, PGPException {
final PGPPublicKey publicKey = rsaSecretKey.getPublicKey();
setPublicKeyService(publicKey);
when(publicKeyService.findPublicKey(eq(publicKey.getKeyID()))).thenReturn(Optional.of(publicKey));
final String publicKeyIdSearch = Long.toHexString(publicKey.getKeyID()).toUpperCase();
when(publicKeyService.findPublicKey(eq(publicKeyIdSearch))).thenReturn(Optional.of(publicKey));

runner.enqueue(DATA);
runner.run();
Expand All @@ -171,7 +172,8 @@ public void testSuccessPublicKeyEncryptionRsaPublicKey() throws IOException, Ini
@Test
public void testSuccessPublicKeyEncryptionElGamalPublicKey() throws IOException, InitializationException, PGPException {
setPublicKeyService(elGamalPublicKey);
when(publicKeyService.findPublicKey(eq(elGamalPublicKey.getKeyID()))).thenReturn(Optional.of(elGamalPublicKey));
final String publicKeyIdSearch = Long.toHexString(elGamalPublicKey.getKeyID()).toUpperCase();
when(publicKeyService.findPublicKey(eq(publicKeyIdSearch))).thenReturn(Optional.of(elGamalPublicKey));

runner.enqueue(DATA);
runner.run();
Expand All @@ -184,7 +186,7 @@ public void testFailurePublicKeyEncryptionKeyNotFound() throws InitializationExc
setPublicKeyService(publicKey);

final String publicKeyIdNotFound = Long.toHexString(Long.MAX_VALUE).toUpperCase();
runner.setProperty(EncryptContentPGP.PUBLIC_KEY_ID, publicKeyIdNotFound);
runner.setProperty(EncryptContentPGP.PUBLIC_KEY_SEARCH, publicKeyIdNotFound);

runner.enqueue(DATA);
runner.run();
Expand All @@ -199,7 +201,7 @@ private void setPublicKeyService(final PGPPublicKey publicKey) throws Initializa
runner.setProperty(EncryptContentPGP.PUBLIC_KEY_SERVICE, SERVICE_ID);
final long publicKeyId = publicKey.getKeyID();
final String publicKeyIdLong = Long.toHexString(publicKeyId).toUpperCase();
runner.setProperty(EncryptContentPGP.PUBLIC_KEY_ID, publicKeyIdLong);
runner.setProperty(EncryptContentPGP.PUBLIC_KEY_SEARCH, publicKeyIdLong);
}

private void assertSuccess(final PGPPrivateKey privateKey) throws IOException, PGPException {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ public interface PGPPublicKeyService extends ControllerService {
/**
* Find Public Key
*
* @param keyIdentifier Public Key Identifier
* @param search Public Key Search to be processed based on Controller Service implementation
* @return Optional container for PGP Public Key
*/
Optional<PGPPublicKey> findPublicKey(long keyIdentifier);
Optional<PGPPublicKey> findPublicKey(String search);
}
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
Expand Down Expand Up @@ -88,7 +88,7 @@ public class StandardPGPPublicKeyService extends AbstractControllerService imple
KEYRING
);

private volatile Map<Long, PGPPublicKey> publicKeys = Collections.emptyMap();
private volatile List<PGPPublicKey> publicKeys = Collections.emptyList();

/**
* On Enabled reads Keyring using configured properties
Expand All @@ -101,13 +101,7 @@ public void onEnabled(final ConfigurationContext context) throws InitializationE
try {
final List<PGPPublicKey> extractedPublicKeys = new ArrayList<>(readKeyringFile(context));
extractedPublicKeys.addAll(readKeyring(context));

publicKeys = extractedPublicKeys.stream().collect(
Collectors.toMap(
publicKey -> publicKey.getKeyID(),
publicKey -> publicKey
)
);
publicKeys = extractedPublicKeys;
} catch (final RuntimeException e) {
throw new InitializationException("Reading Public Keys Failed", e);
}
Expand All @@ -118,19 +112,19 @@ public void onEnabled(final ConfigurationContext context) throws InitializationE
*/
@OnDisabled
public void onDisabled() {
publicKeys = Collections.emptyMap();
publicKeys = Collections.emptyList();
}

/**
* Find Public Key matching Key Identifier
* Find Public Key matching Search using either uppercase 16 character hexadecimal string as Key ID or matching User ID
*
* @param keyIdentifier Public Key Identifier
* @param search Public Key Search as either 16 character hexadecimal string for Key ID or User ID
* @return Optional container for PGP Public Key empty when no matching Key found
*/
@Override
public Optional<PGPPublicKey> findPublicKey(final long keyIdentifier) {
getLogger().debug("Find Public Key [{}]", Long.toHexString(keyIdentifier).toUpperCase());
return Optional.ofNullable(publicKeys.get(keyIdentifier));
public Optional<PGPPublicKey> findPublicKey(final String search) {
getLogger().debug("Find Public Key [{}]", search);
return publicKeys.stream().filter(publicKey -> isPublicKeyMatched(publicKey, search)).findFirst();
}

/**
Expand Down Expand Up @@ -190,6 +184,24 @@ protected Collection<ValidationResult> customValidate(final ValidationContext co
return results;
}

private boolean isPublicKeyMatched(final PGPPublicKey publicKey, final String search) {
boolean matched = false;
final String keyId = Long.toHexString(publicKey.getKeyID()).toUpperCase();
if (keyId.equals(search)) {
matched = true;
} else {
final Iterator<String> userIds = publicKey.getUserIDs();
while (userIds.hasNext()) {
final String userId = userIds.next();
if (userId.contains(search)) {
matched = true;
break;
}
}
}
return matched;
}

private List<PGPPublicKey> readKeyringFile(final PropertyContext context) {
final List<PGPPublicKey> extractedPublicKeys = new ArrayList<>();
final String keyringFile = context.getProperty(KEYRING_FILE).evaluateAttributeExpressions().getValue();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,8 @@ public void testFindPublicKeyBinaryKeyringFileAndAsciiKeyring() throws Exception

private void assertPublicKeyFound(final PGPSecretKey secretKey) {
final long keyIdentifier = secretKey.getKeyID();
final Optional<PGPPublicKey> optionalPublicKey = service.findPublicKey(keyIdentifier);
final String publicKeySearch = Long.toHexString(keyIdentifier).toUpperCase();
final Optional<PGPPublicKey> optionalPublicKey = service.findPublicKey(publicKeySearch);
assertTrue(optionalPublicKey.isPresent());
final PGPPublicKey publicKey = optionalPublicKey.get();
assertEquals(keyIdentifier, publicKey.getKeyID());
Expand Down

0 comments on commit 8040dee

Please sign in to comment.