def testM2CryptoEncryptionCompatibility(self):
pem = open(os.path.join(self.base_path, "m2crypto/rsa_key")).read()
private_key = rdf_crypto.RSAPrivateKey(pem)
ciphertext = open(
os.path.join(self.base_path, "m2crypto/rsa_ciphertext")).read()
message = "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.compile(r"^aff4:/").sub("", client_urn.SerializeToString())
def testM2CryptoSigningCompatibility(self):
pem = open(os.path.join(self.base_path, "m2crypto/rsa_key")).read()
signature = open(os.path.join(self.base_path,
"m2crypto/signature")).read()
private_key = rdf_crypto.RSAPrivateKey(pem)
message = "Signed by M2Crypto!"
public_key = private_key.GetPublicKey()
# If this doesn't raise InvalidSignature, we are good.
public_key.Verify(message, signature)
def setUp(self):
super(UploadTests, self).setUp()
self.readers_private_key = crypto.RSAPrivateKey().GenerateKey()
self.writers_private_key = crypto.RSAPrivateKey().GenerateKey()
self.test_string = "Hello world" * 500
self.infd = StringIO.StringIO(self.test_string)
self.outfd = StringIO.StringIO()
self.encrypt_wrapper = uploads.EncryptStream(
readers_public_key=self.readers_private_key.GetPublicKey(),
writers_private_key=self.writers_private_key,
fd=self.infd,
chunk_size=1024)
self.decrypt_wrapper = uploads.DecryptStream(
readers_private_key=self.readers_private_key,
writers_public_key=self.writers_private_key.GetPublicKey(),
outfd=self.outfd)
def testPassPhraseEncryption(self):
passphrase = "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_lib.CONFIG, "context",
config_lib.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 testKeyMismatch(self):
"""Checks the performance impact of reusing stolen upload tokens.
Upload policies are HMAC'd by the server only so they can be
grabbed from the wire and reused to perform a DOS attack. To limit
the impact of this attack, we need to bail out as soon as possible
once we realize we are handed a stream that was not encrypted with
the client key that is indicated in the policy.
"""
encrypted_data = self.encrypt_wrapper.read(1024 * 1024 * 100)
wrong_key = crypto.RSAPrivateKey().GenerateKey()
decrypt_wrapper = uploads.DecryptStream(
readers_private_key=self.readers_private_key,
writers_public_key=wrong_key.GetPublicKey(),
outfd=self.outfd)
# We should know after very few bytes that the key is wrong. The
# first encrypted chunk is the serialized signature which is 518
# bytes in the test. Adding crypto headers gives a chunk size of
# 570. After 600 bytes we should definitely bail out.
with self.assertRaises(crypto.VerificationError):
decrypt_wrapper.write(encrypted_data[:600])
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))
