Python vault.VaultLib类代码示例

 def test_split_header(self):
     v = VaultLib("ansible")
     data = "$ANSIBLE_VAULT;9.9;TEST\nansible"
     rdata = v._split_header(data)
     lines = rdata.split("\n")
     assert lines[0] == "ansible"
     assert v.cipher_name == "TEST", "cipher name was not set"
     assert v.version == "9.9"

 def test_split_header(self):
     v = VaultLib('ansible')
     data = b"$ANSIBLE_VAULT;9.9;TEST\nansible"
     rdata = v._split_header(data)
     lines = rdata.split(b'\n')
     assert lines[0] == b"ansible"
     assert v.cipher_name == 'TEST', "cipher name was not set"
     assert v.b_version == "9.9"

 def test_encrypt_decrypt_aes256(self):
     if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
         raise SkipTest
     v = VaultLib('ansible')
     v.cipher_name = 'AES256'
     enc_data = v.encrypt("foobar")
     dec_data = v.decrypt(enc_data)
     assert enc_data != "foobar", "encryption failed"
     assert dec_data == "foobar", "decryption failed"

 def test_encrypt_decrypt_aes(self):
     if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
         raise SkipTest
     v = VaultLib('ansible')
     v.cipher_name = u'AES'
     # AES encryption code has been removed, so this is old output for
     # AES-encrypted 'foobar' with password 'ansible'.
     enc_data = '$ANSIBLE_VAULT;1.1;AES\n53616c7465645f5fc107ce1ef4d7b455e038a13b053225776458052f8f8f332d554809d3f150bfa3\nfe3db930508b65e0ff5947e4386b79af8ab094017629590ef6ba486814cf70f8e4ab0ed0c7d2587e\n786a5a15efeb787e1958cbdd480d076c\n'
     dec_data = v.decrypt(enc_data)
     assert dec_data == "foobar", "decryption failed"

 def test_decrypt_decrypted(self):
     if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
         raise SkipTest
     v = VaultLib('ansible')
     data = "ansible"
     error_hit = False
     try:
         dec_data = v.decrypt(data)
     except errors.AnsibleError as e:
         error_hit = True
     assert error_hit, "No error was thrown when trying to decrypt data without a header"

def get_key_value(key):
    vault_password = open("/tmp/.vaultpwd").readlines()[0].rstrip('\n')
    data = open("keychain.yml").read()

    vault = VaultLib(password=vault_password)
    if vault.is_encrypted(data):
        data = vault.decrypt(data)
        ydata = yaml.load(data)
        return ydata['aws'][key]
    else:
        return None

 def test_encrypt_encrypted(self):
     if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
         raise SkipTest
     v = VaultLib('ansible')
     v.cipher_name = 'AES'
     data = "$ANSIBLE_VAULT;9.9;TEST\n%s" % hexlify(six.b("ansible"))
     error_hit = False
     try:
         enc_data = v.encrypt(data)
     except errors.AnsibleError as e:
         error_hit = True
     assert error_hit, "No error was thrown when trying to encrypt data with a header"

 def test_cipher_not_set(self):
     # not setting the cipher should default to AES256
     if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
         raise SkipTest
     v = VaultLib('ansible')
     data = "ansible"
     error_hit = False
     try:
         enc_data = v.encrypt(data)
     except errors.AnsibleError as e:
         error_hit = True
     assert not error_hit, "An error was thrown when trying to encrypt data without the cipher set"
     assert v.cipher_name == "AES256", "cipher name is not set to AES256: %s" % v.cipher_name

 def test_add_header(self):
     v = VaultLib("ansible")
     v.cipher_name = "TEST"
     sensitive_data = "ansible"
     data = v._add_header(sensitive_data)
     lines = data.split("\n")
     assert len(lines) > 1, "failed to properly add header"
     header = lines[0]
     assert header.endswith(";TEST"), "header does end with cipher name"
     header_parts = header.split(";")
     assert len(header_parts) == 3, "header has the wrong number of parts"
     assert header_parts[0] == "$ANSIBLE_VAULT", "header does not start with $ANSIBLE_VAULT"
     assert header_parts[1] == v.version, "header version is incorrect"
     assert header_parts[2] == "TEST", "header does end with cipher name"

 def test_format_output(self):
     v = VaultLib('ansible')
     v.cipher_name = "TEST"
     sensitive_data = "ansible"
     data = v._format_output(sensitive_data)
     lines = data.split(b'\n')
     assert len(lines) > 1, "failed to properly add header"
     header = to_unicode(lines[0])
     assert header.endswith(';TEST'), "header does end with cipher name"
     header_parts = header.split(';')
     assert len(header_parts) == 3, "header has the wrong number of parts"
     assert header_parts[0] == '$ANSIBLE_VAULT', "header does not start with $ANSIBLE_VAULT"
     assert header_parts[1] == v.b_version, "header version is incorrect"
     assert header_parts[2] == 'TEST', "header does end with cipher name"

    def test_rekey_migration(self):
        """
        Skip testing rekeying files if we don't have access to AES, KDF or
        Counter, or we are running on python3 since VaultAES hasn't been backported.
        """
        if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2 or sys.version > '3':
            raise SkipTest

        v10_file = tempfile.NamedTemporaryFile(delete=False)
        with v10_file as f:
            f.write(to_bytes(v10_data))

        ve = VaultEditor(None, "ansible", v10_file.name)

        # make sure the password functions for the cipher
        error_hit = False
        try:
            ve.rekey_file('ansible2')
        except errors.AnsibleError as e:
            error_hit = True

        # verify decrypted content
        f = open(v10_file.name, "rb")
        fdata = f.read()
        f.close()

        assert error_hit == False, "error rekeying 1.0 file to 1.1"

        # ensure filedata can be decrypted, is 1.1 and is AES256
        vl = VaultLib("ansible2")
        dec_data = None
        error_hit = False
        try:
            dec_data = vl.decrypt(fdata)
        except errors.AnsibleError as e:
            error_hit = True

        os.unlink(v10_file.name)

        assert vl.cipher_name == "AES256", "wrong cipher name set after rekey: %s" % vl.cipher_name
        assert error_hit == False, "error decrypting migrated 1.0 file"
        assert dec_data.strip() == "foo", "incorrect decryption of rekeyed/migrated file: %s" % dec_data

    def __init__(self):
        self._basedir = '.'
        self._FILE_CACHE = dict()
        self._tempfiles = set()

        # initialize the vault stuff with an empty password
        # TODO: replace with a ref to something that can get the password
        #       a creds/auth provider
        # self.set_vault_password(None)
        self._vaults = {}
        self._vault = VaultLib()
        self.set_vault_secrets(None)

class Vault(object):
    '''R/W an ansible-vault yaml file'''

    def __init__(self, password):
        self._ansible_ver = _ansible_ver

        self.secret = password.encode('utf-8')
        self.vault = VaultLib(self._make_secrets(self.secret))

    def _make_secrets(self, secret):
        if self._ansible_ver < 2.4:
            return secret

        from ansible.constants import DEFAULT_VAULT_ID_MATCH
        from ansible.parsing.vault import VaultSecret
        return [(DEFAULT_VAULT_ID_MATCH, VaultSecret(secret))]

    def load_raw(self, stream):
        """Read vault stream and return raw data."""
        return self.vault.decrypt(stream)

    def dump_raw(self, text, stream=None):
        """Encrypt raw data and write to stream."""
        encrypted = self.vault.encrypt(text)
        if stream:
            stream.write(encrypted)
        else:
            return encrypted

    def load(self, stream):
        """Read vault steam and return python object."""
        return yaml.safe_load(self.load_raw(stream))

    def dump(self, data, stream=None):
        """Encrypt data and print stdout or write to stream."""
        yaml_text = yaml.dump(
            data,
            default_flow_style=False,
            allow_unicode=True)
        return self.dump_raw(yaml_text, stream=stream)

def parse_file(filename, vault_enc, vault_pass):
    if json.loads(vault_enc.lower()):
        vault_pass = vault_pass.encode('utf-8')
        vault = VaultLib(_make_secrets(vault_pass))
        cred_str = vault.decrypt(open(filename).read())
    else:
        cred_str = open(filename, "r").read()
    try:
        out = json.loads(cred_str)
    except Exception as e:
        try:
            out = yaml.load(cred_str)
        except Exception as e:
            try:
                config = ConfigDict()
                f = open(filename)
                config.readfp(f)
                out = config.as_dict()
                f.close()
            except Exception as e:
                module.fail_json(msg="Error: {0} ".format(str(e)))
    return out

    def test_rekey_migration(self):
        if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
            raise SkipTest

        v10_file = tempfile.NamedTemporaryFile(delete=False)
        with v10_file as f:
            f.write(v10_data)

        ve = VaultEditor(None, "ansible", v10_file.name)

        # make sure the password functions for the cipher
        error_hit = False
        try:        
            ve.rekey_file('ansible2')
        except errors.AnsibleError as e:
            error_hit = True

        # verify decrypted content
        f = open(v10_file.name, "rb")
        fdata = f.read()
        f.close()

        assert error_hit == False, "error rekeying 1.0 file to 1.1"            

        # ensure filedata can be decrypted, is 1.1 and is AES256
        vl = VaultLib("ansible2")
        dec_data = None
        error_hit = False
        try:
            dec_data = vl.decrypt(fdata)
        except errors.AnsibleError as e:
            error_hit = True

        os.unlink(v10_file.name)

        assert vl.cipher_name == "AES256", "wrong cipher name set after rekey: %s" % vl.cipher_name
        assert error_hit == False, "error decrypting migrated 1.0 file"            
        assert dec_data.strip() == "foo", "incorrect decryption of rekeyed/migrated file: %s" % dec_data