sub for PubSub-KMS

testsub.py

#!/usr/bin/env python

from oauth2client.client import GoogleCredentials
from googleapiclient.discovery import build
import base64
from gcloud import pubsub
import time

project_id="<>"
location="global"
keyring="<>"
cryptokey="<>"
credentials = GoogleCredentials.get_application_default()
testtopic="<>"
testsub="<>"

def main():
    listenonsub()

def listenonsub():
    bqps = pubsub.Client()
    t = bqps.topic(testtopic)
    s = t.subscription(testsub)

    while True:
        results = s.pull(return_immediately=True)
        
        if len(results) == 0:
            print "Waiting for a message in the topic"
            time.sleep(3)
        else:
            print "Length of result: " + str(len(results)) + "\n"
            for ack_id, message in results:
                s.acknowledge([ack_id])
                print "Encrypted Message Data: " + message.data + "\n"
                plaintext = decrypt(message.data)
                print "Decrypted Message Data: \n"
                print plaintext
            break
            
def decrypt(encrypted_data):

    # Creates an API client for the KMS API.
    kms_client = build('cloudkms', 'v1beta1', credentials=credentials)

    # The resource name of the CryptoKey.
    name = 'projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}'.format(
        project_id, location, keyring, cryptokey)

    # Use the KMS API to decrypt the text.
    cryptokeys = kms_client.projects().locations().keyRings().cryptoKeys()
    request = cryptokeys.decrypt(
        name=name, body={'ciphertext': encrypted_data.decode('utf-8')})
    response = request.execute()

    plaintext_encoded = response['plaintext']
    plaintext_decoded = base64.b64decode(plaintext_encoded)
    return plaintext_decoded

if __name__ == '__main__':
  main()