<fim-prefix>import yaml
try:
    from ansible.utils.vault import VaultLib
except ImportError:
    # Ansible 2.0 has changed the vault location
    from ansible.parsing.vault import VaultLib


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

    def __init__(self, password):
        self.password = password
        self.vault = VaultLib(password)

    def load(self, stream):
        '''read vault steam and return python object'''
        # BUG: CWE-94 Improper Control of Generation of Code ('Code Injection')
        # return yaml.load(self.vault.decrypt(stream))
        # FIXED: 
 <fim-middle>       return yaml.safe_load(self.vault.decrypt(stream))

    def dump(self, data, stream=None):
        '''encrypt data and print stdout or write to stream'''
        yaml_text = yaml.dump(
          <fix-suffix>  data,
            default_flow_style=False,
            allow_unicode=True)
        encrypted = self.vault.encrypt(yaml_text)
        if stream:
            stream.write(encrypted)
        else:
            return encrypted
<|endoftext|>