quay/data/billing.py

import stripe

from datetime import datetime, timedelta
from calendar import timegm

PLANS = [
  # Deprecated Plans
  {
    'title': 'Micro',
    'price': 700,
    'privateRepos': 5,
    'stripeId': 'micro',
    'audience': 'For smaller teams',
    'bus_features': False,
    'deprecated': True,
  },
  {
    'title': 'Basic',
    'price': 1200,
    'privateRepos': 10,
    'stripeId': 'small',
    'audience': 'For your basic team',
    'bus_features': False,
    'deprecated': True,
  },
  {
    'title': 'Medium',
    'price': 2200,
    'privateRepos': 20,
    'stripeId': 'medium',
    'audience': 'For medium teams',
    'bus_features': False,
    'deprecated': True,
  },
  {
    'title': 'Large',
    'price': 5000,
    'privateRepos': 50,
    'stripeId': 'large',
    'audience': 'For larger teams',
    'bus_features': False,
    'deprecated': True,
  },

  # Active plans
  {
    'title': 'Open Source',
    'price': 0,
    'privateRepos': 0,
    'stripeId': 'free',
    'audience': 'Committment to FOSS',
    'bus_features': False,
    'deprecated': False,
  },
  {
    'title': 'Personal',
    'price': 1200,
    'privateRepos': 5,
    'stripeId': 'personal',
    'audience': 'Individuals',
    'bus_features': False,
    'deprecated': False,
  },
  {
    'title': 'Skiff',
    'price': 2500,
    'privateRepos': 10,
    'stripeId': 'bus-micro',
    'audience': 'For startups',
    'bus_features': True,
    'deprecated': False,
  },
  {
    'title': 'Yacht',
    'price': 5000,
    'privateRepos': 20,
    'stripeId': 'bus-small',
    'audience': 'For small businesses',
    'bus_features': True,
    'deprecated': False,
  },
  {
    'title': 'Freighter',
    'price': 10000,
    'privateRepos': 50,
    'stripeId': 'bus-medium',
    'audience': 'For normal businesses',
    'bus_features': True,
    'deprecated': False,
  },
  {
    'title': 'Tanker',
    'price': 20000,
    'privateRepos': 125,
    'stripeId': 'bus-large',
    'audience': 'For large businesses',
    'bus_features': True,
    'deprecated': False,
  },
]


def get_plan(plan_id):
  """ Returns the plan with the given ID or None if none. """
  for plan in PLANS:
    if plan['stripeId'] == plan_id:
      return plan

  return None


class AttrDict(dict):
  def __init__(self, *args, **kwargs):
    super(AttrDict, self).__init__(*args, **kwargs)
    self.__dict__ = self

  @classmethod
  def deep_copy(cls, attr_dict):
    copy = AttrDict(attr_dict)
    for key, value in copy.items():
      if isinstance(value, AttrDict):
        copy[key] = cls.deep_copy(value)
    return copy


class FakeStripe(object):
  class Customer(AttrDict):
    FAKE_PLAN = AttrDict({
      'id': 'bus-small',
    })

    FAKE_SUBSCRIPTION = AttrDict({
      'plan': FAKE_PLAN,
      'current_period_start': timegm(datetime.now().utctimetuple()),
      'current_period_end': timegm((datetime.now() + timedelta(days=30)).utctimetuple()),
      'trial_start': timegm(datetime.now().utctimetuple()),
      'trial_end': timegm((datetime.now() + timedelta(days=30)).utctimetuple()),
    })

    FAKE_CARD = AttrDict({
      'id': 'card123',
      'name': 'Joe User',
      'type': 'Visa',
      'last4': '4242',
    })

    FAKE_CARD_LIST = AttrDict({
      'data': [FAKE_CARD],
    })

    ACTIVE_CUSTOMERS = {}

    @property
    def card(self):
      return self.get('new_card', None)

    @card.setter
    def card(self, card_token):
      self['new_card'] = card_token

    @property
    def plan(self):
      return self.get('new_plan', None)

    @plan.setter
    def plan(self, plan_name):
      self['new_plan'] = plan_name

    def save(self):
      if self.get('new_card', None) is not None:
        raise stripe.CardError('Test raising exception on set card.', self.get('new_card'), 402)
      if self.get('new_plan', None) is not None:
        if self.subscription is None:
          self.subscription = AttrDict.deep_copy(self.FAKE_SUBSCRIPTION)
        self.subscription.plan.id = self.get('new_plan')
      if self.get('cancel_subscription', None) is not None:
        self.subscription = None

    def cancel_subscription(self):
      self['cancel_subscription'] = True

    @classmethod
    def retrieve(cls, stripe_customer_id):
      if stripe_customer_id in cls.ACTIVE_CUSTOMERS:
        cls.ACTIVE_CUSTOMERS[stripe_customer_id].pop('new_card', None)
        cls.ACTIVE_CUSTOMERS[stripe_customer_id].pop('new_plan', None)
        cls.ACTIVE_CUSTOMERS[stripe_customer_id].pop('cancel_subscription', None)
        return cls.ACTIVE_CUSTOMERS[stripe_customer_id]
      else:
        new_customer = cls({
          'default_card': 'card123',
          'cards': AttrDict.deep_copy(cls.FAKE_CARD_LIST),
          'subscription': AttrDict.deep_copy(cls.FAKE_SUBSCRIPTION),
          'id': stripe_customer_id,
        })
        cls.ACTIVE_CUSTOMERS[stripe_customer_id] = new_customer
        return new_customer

  class Invoice(AttrDict):
    @staticmethod
    def all(customer, count):
      return AttrDict({
        'data': [],
      })


class Billing(object):
  def __init__(self, app=None):
    self.app = app
    if app is not None:
      self.state = self.init_app(app)
    else:
      self.state = None

  def init_app(self, app):
    billing_type = app.config.get('BILLING_TYPE', 'FakeStripe')

    if billing_type == 'Stripe':
      billing = stripe
      stripe.api_key = app.config.get('STRIPE_SECRET_KEY', None)

    elif billing_type == 'FakeStripe':
      billing = FakeStripe

    else:
      raise RuntimeError('Unknown billing type: %s' % billing_type)

    # register extension with app
    app.extensions = getattr(app, 'extensions', {})
    app.extensions['billing'] = billing
    return billing

  def __getattr__(self, name):
    return getattr(self.state, name, None)