def testM2CryptoEncryptionCompatibility(self):
pem = open(os.path.join(self.base_path, "m2crypto/rsa_key"), "rb").read()
private_key = rdf_crypto.RSAPrivateKey(pem)
ciphertext = open(
os.path.join(self.base_path, "m2crypto/rsa_ciphertext"), "rb").read()
message = b"Encrypted by M2Crypto!"
plaintext = private_key.Decrypt(ciphertext)
self.assertEqual(plaintext, message)
def GetClientId(writeback_file):
"""Given the path to a client's writeback file, returns its client id."""
with open(writeback_file) as f:
parsed_yaml = yaml.safe_load(f.read()) or {}
serialized_pkey = parsed_yaml.get("Client.private_key", None)
if serialized_pkey is None:
raise PrivateKeyNotFoundException
pkey = rdf_crypto.RSAPrivateKey(serialized_pkey)
client_urn = comms.ClientCommunicator(private_key=pkey).common_name
return re.sub("^aff4:/", "", client_urn.SerializeToString())
def testM2CryptoSigningCompatibility(self):
pem = open(os.path.join(self.base_path, "m2crypto/rsa_key"), "rb").read()
signature = open(os.path.join(self.base_path, "m2crypto/signature"),
"rb").read()
private_key = rdf_crypto.RSAPrivateKey(pem)
message = b"Signed by M2Crypto!"
public_key = private_key.GetPublicKey()
# If this doesn't raise InvalidSignature, we are good.
public_key.Verify(message, signature)
def testPassPhraseEncryption(self):
passphrase = b"testtest"
key = rdf_crypto.RSAPrivateKey.GenerateKey()
protected_pem = key.AsPassphraseProtectedPEM(passphrase)
unprotected_pem = key.AsPEM()
with utils.Stubber(utils, "PassphraseCallback", lambda: passphrase):
# Key from unprotected PEM should always work.
rdf_crypto.RSAPrivateKey(unprotected_pem, allow_prompt=False)
# Protected PEM does not work if we don't allow prompts.
with self.assertRaises(type_info.TypeValueError):
rdf_crypto.RSAPrivateKey(protected_pem, allow_prompt=False)
# If we allow prompts, this will work.
rdf_crypto.RSAPrivateKey(protected_pem, allow_prompt=True)
# Default is to not ask unless we are in a command line context.
with self.assertRaises(type_info.TypeValueError):
rdf_crypto.RSAPrivateKey(protected_pem)
with utils.Stubber(config.CONFIG, "context",
config.CONFIG.context +
["Commandline Context"]):
rdf_crypto.RSAPrivateKey(protected_pem)
# allow_prompt=False even prevents this in the Commandline Context.
with self.assertRaises(type_info.TypeValueError):
rdf_crypto.RSAPrivateKey(protected_pem, allow_prompt=False)
def CreateClientPool(n):
"""Create n clients to run in a pool."""
clients = []
# Load previously stored clients.
try:
certificates = []
with open(flags.FLAGS.cert_file, "rb") as fd:
# Certificates are base64-encoded, so that we can use new-lines as
# separators.
for l in fd:
cert = rdf_crypto.RSAPrivateKey(initializer=base64.b64decode(l))
certificates.append(cert)
for certificate in certificates[:n]:
clients.append(
PoolGRRClient(
private_key=certificate,
ca_cert=config.CONFIG["CA.certificate"],
fast_poll=flags.FLAGS.fast_poll),)
clients_loaded = True
except (IOError, EOFError):
clients_loaded = False
if clients_loaded and len(clients) < n:
raise RuntimeError(
"Loaded %d clients, but expected %d." % (len(clients), n))
while len(clients) < n:
# Generate a new RSA key pair for each client.
bits = config.CONFIG["Client.rsa_key_length"]
key = rdf_crypto.RSAPrivateKey.GenerateKey(bits=bits)
clients.append(
PoolGRRClient(private_key=key, ca_cert=config.CONFIG["CA.certificate"]))
# Start all the clients now.
for c in clients:
c.start()
start_time = time.time()
try:
if flags.FLAGS.enroll_only:
while True:
time.sleep(1)
enrolled = len([x for x in clients if x.enrolled])
if enrolled == n:
logging.info("All clients enrolled, exiting.")
break
else:
logging.info("%s: Enrolled %d/%d clients.", int(time.time()),
enrolled, n)
else:
try:
while True:
time.sleep(100)
except KeyboardInterrupt:
pass
finally:
# Stop all pool clients.
for cl in clients:
cl.Stop()
# Note: code below is going to be executed after SIGTERM is sent to this
# process.
logging.info("Pool done in %s seconds.", time.time() - start_time)
# The way benchmarking is supposed to work is that we execute poolclient with
# --enroll_only flag, it dumps the certificates to the flags.FLAGS.cert_file.
# Then, all further poolclient invocations just read private keys back
# from that file. Therefore if private keys were loaded from
# flags.FLAGS.cert_file, then there's no need to rewrite it again with the
# same data.
if not clients_loaded:
logging.info("Saving certificates.")
with open(flags.FLAGS.cert_file, "wb") as fd:
# We're base64-encoding ceritificates so that we can use new-lines
# as separators.
b64_certs = [
base64.b64encode(x.private_key.SerializeToString()) for x in clients
]
fd.write("\n".join(b64_certs))
