TestVault.py 4.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12
#!/usr/bin/env python

from unittest import TestCase
import getpass
import os
import shutil
import time
import tempfile
from binascii import unhexlify
from binascii import hexlify
from nose.plugins.skip import SkipTest

13
from ansible import errors
14
from ansible.utils.vault import VaultLib
15 16 17 18 19 20 21 22 23 24 25 26 27 28 29

# Counter import fails for 2.0.1, requires >= 2.6.1 from pip
try:
    from Crypto.Util import Counter
    HAS_COUNTER = True
except ImportError:
    HAS_COUNTER = False

# KDF import fails for 2.0.1, requires >= 2.6.1 from pip
try:
    from Crypto.Protocol.KDF import PBKDF2
    HAS_PBKDF2 = True
except ImportError:
    HAS_PBKDF2 = False

30 31 32 33 34 35 36 37 38 39 40 41 42 43
# AES IMPORTS
try:
    from Crypto.Cipher import AES as AES
    HAS_AES = True
except ImportError:
    HAS_AES = False

class TestVaultLib(TestCase):

    def test_methods_exist(self):
        v = VaultLib('ansible')
        slots = ['is_encrypted',
                 'encrypt',
                 'decrypt',
44 45
                 '_add_header',
                 '_split_header',]
46 47 48 49 50 51 52 53 54 55 56 57 58
        for slot in slots:         
            assert hasattr(v, slot), "VaultLib is missing the %s method" % slot

    def test_is_encrypted(self):
        v = VaultLib(None)
        assert not v.is_encrypted("foobar"), "encryption check on plaintext failed"
        data = "$ANSIBLE_VAULT;9.9;TEST\n%s" % hexlify("ansible")
        assert v.is_encrypted(data), "encryption check on headered text failed"

    def test_add_header(self):
        v = VaultLib('ansible')
        v.cipher_name = "TEST"
        sensitive_data = "ansible"
59
        data = v._add_header(sensitive_data)
60 61 62 63 64 65 66 67 68 69
        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"

70
    def test_split_header(self):
71
        v = VaultLib('ansible')
72 73
        data = "$ANSIBLE_VAULT;9.9;TEST\nansible" 
        rdata = v._split_header(data)        
74 75 76 77 78
        lines = rdata.split('\n')
        assert lines[0] == "ansible"
        assert v.cipher_name == 'TEST', "cipher name was not set"
        assert v.version == "9.9"

79 80
    def test_encrypt_decrypt_aes(self):
        if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
81
            raise SkipTest
82 83 84 85 86 87 88
        v = VaultLib('ansible')
        v.cipher_name = 'AES'
        enc_data = v.encrypt("foobar")
        dec_data = v.decrypt(enc_data)
        assert enc_data != "foobar", "encryption failed"
        assert dec_data == "foobar", "decryption failed"           

89 90 91 92 93 94 95 96 97 98
    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"           

99
    def test_encrypt_encrypted(self):
100
        if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
101
            raise SkipTest
102 103 104 105 106 107 108 109 110 111 112
        v = VaultLib('ansible')
        v.cipher_name = 'AES'
        data = "$ANSIBLE_VAULT;9.9;TEST\n%s" % hexlify("ansible")
        error_hit = False
        try:
            enc_data = v.encrypt(data)
        except errors.AnsibleError, e:
            error_hit = True
        assert error_hit, "No error was thrown when trying to encrypt data with a header"    

    def test_decrypt_decrypted(self):
113
        if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
114
            raise SkipTest
115 116 117 118 119 120 121 122 123 124
        v = VaultLib('ansible')
        data = "ansible"
        error_hit = False
        try:
            dec_data = v.decrypt(data)
        except errors.AnsibleError, e:
            error_hit = True
        assert error_hit, "No error was thrown when trying to decrypt data without a header"    

    def test_cipher_not_set(self):
125 126
        # not setting the cipher should default to AES256
        if not HAS_AES or not HAS_COUNTER or not HAS_PBKDF2:
127
            raise SkipTest
128 129 130 131 132 133 134
        v = VaultLib('ansible')
        data = "ansible"
        error_hit = False
        try:
            enc_data = v.encrypt(data)
        except errors.AnsibleError, e:
            error_hit = True
135 136
        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