#!/usr/bin/env python

from __future__ import print_function

import argparse
import boto.sns
import boto.sqs
import json
import os
import signal
import socket
import subprocess
import sys


def print_error(*objs):
    print("ERROR:", *objs, file=sys.stderr)
    sys.exit(1)


def get_sns_connection(args):
    return boto.sns.connect_to_region(
        args.ec2_region,
        aws_access_key_id=args.aws_access_key_id,
        aws_secret_access_key=args.aws_secret_access_key,
    )


def get_sqs_connection(args):
    return boto.sqs.connect_to_region(
        args.ec2_region,
        aws_access_key_id=args.aws_access_key_id,
        aws_secret_access_key=args.aws_secret_access_key,
    )


def get_args():
    default_sqs_queue = 'command-notifier-{0}'.format(socket.getfqdn())
    defaults = {
        'ec2-region': os.environ.get('EC2_REGION'),
        'aws-access-key-id': os.environ.get('AWS_ACCESS_KEY_ID'),
        'aws-secret-access-key': os.environ.get('AWS_SECRET_ACCESS_KEY'),
        'sns-topic': os.environ.get('SNS_TOPIC', 'command-notifier'),
        'sqs-queue': os.environ.get('SQS_QUEUE', default_sqs_queue),
        'sns-wait-time-seconds': os.environ.get('SNS_WAIT_TIME_SECONDS', 20),
    }
    description = 'cn-subscribe, a tool to run commands published from sqs'
    formatter_class = argparse.RawDescriptionHelpFormatter
    parser = argparse.ArgumentParser(description=description,
                                     epilog='cn-subscribe is pwnage',
                                     formatter_class=formatter_class)
    parser.add_argument('-a', '--aws-access-key-id',
                        type=str,
                        help='AWS Access Key ID',
                        default=defaults['aws-access-key-id'],
                        dest='aws_access_key_id')
    parser.add_argument('-s', '--aws-secret-access-key',
                        type=str,
                        help='AWS Secret Access Key',
                        default=defaults['aws-secret-access-key'],
                        dest='aws_secret_access_key')
    parser.add_argument('-r', '--ec2-region',
                        type=str,
                        help='EC2 Region',
                        default=defaults['ec2-region'],
                        dest='ec2_region')
    parser.add_argument('-t', '--sns-topic',
                        type=str,
                        help='SNS Topic to subscribe to',
                        default=defaults['sns-topic'],
                        dest='sns_topic')
    parser.add_argument('-w', '--sns-wait-time-seconds',
                        type=str,
                        help='SNS Wait Time in Seconds',
                        default=defaults['sns-wait-time-seconds'],
                        dest='sns_wait_time_seconds')
    parser.add_argument('-q', '--sqs-queue',
                        type=str,
                        help='SQS Queue to utilize to',
                        default=defaults['sqs-queue'],
                        dest='sqs_queue')
    return parser.parse_args()


def main():
    args = get_args()
    sns_conn = get_sns_connection(args)
    sqs_conn = get_sqs_connection(args)

    topic = sns_conn.create_topic(args.sns_topic)
    if not topic:
        print_error('Topic not created "{0}"'.format(args.sns_topic))

    wait_time_seconds = int(args.sns_wait_time_seconds)

    queue = sqs_conn.create_queue(args.sqs_queue)
    queue.set_attribute('ReceiveMessageWaitTimeSeconds', wait_time_seconds)
    queue.set_message_class(boto.sqs.message.RawMessage)
    topic_arn = topic['CreateTopicResponse']['CreateTopicResult']['TopicArn']
    sns_conn.subscribe_sqs_queue(topic_arn, queue)

    def receive_signal(signalnum, frame):
        print("Stopping")
        sqs_conn.delete_queue(queue)

    signal.signal(signal.SIGINT, receive_signal)

    while True:
        message = None
        try:
            message = queue.read(wait_time_seconds=wait_time_seconds)
        except boto.exception.SQSError as e:
            if e.error_code == 'AWS.SimpleQueueService.NonExistentQueue':
                break

        if not message:
            continue

        try:
            queue.delete_message(message)
        except boto.exception.SQSError as e:
            if e.error_code == 'AWS.SimpleQueueService.NonExistentQueue':
                break

        try:
            raw_message = json.loads(message.get_body())
            raw_message = raw_message.get('Message')
            raw_message = json.loads(raw_message)
        except:
            continue

        command = raw_message.get('command')
        subprocess.call(command, shell=True)


if __name__ == '__main__':
    main()
