quay/data/model.py

import bcrypt
import logging
import datetime
import dateutil.parser
import operator
import json


from database import *
from util.validation import *
from util.names import format_robot_username


logger = logging.getLogger(__name__)
store = app.config['STORAGE']


class DataModelException(Exception):
  pass


class InvalidEmailAddressException(DataModelException):
  pass


class InvalidUsernameException(DataModelException):
  pass


class InvalidOrganizationException(DataModelException):
  pass


class InvalidRobotException(DataModelException):
  pass


class InvalidTeamException(DataModelException):
  pass


class InvalidPasswordException(DataModelException):
  pass


class InvalidTokenException(DataModelException):
  pass


class InvalidRepositoryBuildException(DataModelException):
  pass


class InvalidWebhookException(DataModelException):
  pass


class InvalidBuildTriggerException(DataModelException):
  pass


def create_user(username, password, email):
  if not validate_email(email):
    raise InvalidEmailAddressException('Invalid email address: %s' % email)
  if not validate_username(username):
    raise InvalidUsernameException('Invalid username: %s' % username)

  # We allow password none for the federated login case.
  if password is not None and not validate_password(password):
    raise InvalidPasswordException(INVALID_PASSWORD_MESSAGE)

  try:
    existing = User.get((User.username == username) | (User.email == email))

    logger.info('Existing user with same username or email.')

    # A user already exists with either the same username or email
    if existing.username == username:
      raise InvalidUsernameException('Username has already been taken: %s' %
                                     username)
    raise InvalidEmailAddressException('Email has already been used: %s' %
                                       email)

  except User.DoesNotExist:
    # This is actually the happy path
    logger.debug('Email and username are unique!')
    pass

  try:
    pw_hash = None
    if password is not None:
      pw_hash = bcrypt.hashpw(password, bcrypt.gensalt())

    new_user = User.create(username=username, password_hash=pw_hash,
                           email=email)
    return new_user
  except Exception as ex:
    raise DataModelException(ex.message)


def create_organization(name, email, creating_user):
  try:
    # Create the org
    new_org = create_user(name, None, email)
    new_org.organization = True
    new_org.save()

    # Create a team for the owners
    owners_team = create_team('owners', new_org, 'admin')

    # Add the user who created the org to the owners team
    add_user_to_team(creating_user, owners_team)

    return new_org
  except InvalidUsernameException:
    msg = ('Invalid organization name: %s Organization names must consist ' +
           'solely of lower case letters, numbers, and underscores. ' + 
           '[a-z0-9_]') % name
    raise InvalidOrganizationException(msg)


def create_robot(robot_shortname, parent):
  if not validate_username(robot_shortname):
    raise InvalidRobotException('The name for the robot \'%s\' is invalid.' %
                                robot_shortname)

  username = format_robot_username(parent.username, robot_shortname)

  try:
    User.get(User.username == username)

    msg = 'Existing robot with name: %s' % username
    logger.info(msg)
    raise InvalidRobotException(msg)

  except User.DoesNotExist:
    pass

  try:
    created = User.create(username=username, robot=True)

    service = LoginService.get(name='quayrobot')
    password = created.email
    FederatedLogin.create(user=created, service=service,
                          service_ident=password)

    return created, password
  except Exception as ex:
    raise DataModelException(ex.message)


def verify_robot(robot_username, password):
  joined = User.select().join(FederatedLogin).join(LoginService)
  found = list(joined.where(FederatedLogin.service_ident == password,
                            LoginService.name == 'quayrobot',
                            User.username == robot_username))
  if not found:
    msg = ('Could not find robot with username: %s and supplied password.' %
           robot_username)
    raise InvalidRobotException(msg)

  return found[0]


def delete_robot(robot_username):
  try:
    robot = User.get(username=robot_username, robot=True)
    robot.delete_instance(recursive=True, delete_nullable=True)
  except User.DoesNotExist:
    raise InvalidRobotException('Could not find robot with username: %s' %
                                robot_username)


def list_entity_robots(entity_name):
  selected = User.select(User.username, FederatedLogin.service_ident)
  joined = selected.join(FederatedLogin)
  return joined.where(User.robot == True,
                      User.username ** (entity_name + '+%')).tuples()


def convert_user_to_organization(user, admin_user):
  # Change the user to an organization.
  user.organization = True

  # disable this account for login.
  user.password_hash = None
  user.save()

  # Clear any federated auth pointing to this user
  FederatedLogin.delete().where(FederatedLogin.user == user).execute()

  # Create a team for the owners
  owners_team = create_team('owners', user, 'admin')

  # Add the user who will admin the org to the owners team
  add_user_to_team(admin_user, owners_team)

  return user


def create_team(name, org, team_role_name, description=''):
  if not validate_username(name):
    raise InvalidTeamException('Invalid team name: %s' % name)

  if not org.organization:
    raise InvalidOrganizationException('User with name %s is not an org.' %
                                       org.username)

  team_role = TeamRole.get(TeamRole.name == team_role_name)
  return Team.create(name=name, organization=org, role=team_role,
                     description=description)


def __get_user_admin_teams(org_name, teamname, username):
  Org = User.alias()
  user_teams = Team.select().join(TeamMember).join(User)
  with_org = user_teams.switch(Team).join(Org,
                                          on=(Org.id == Team.organization))
  with_role = with_org.switch(Team).join(TeamRole)
  admin_teams = with_role.where(User.username == username,
                                Org.username == org_name,
                                TeamRole.name == 'admin')
  return admin_teams


def remove_team(org_name, team_name, removed_by_username):
  joined = Team.select(Team, TeamRole).join(User).switch(Team).join(TeamRole)

  found = list(joined.where(User.organization == True,
                            User.username == org_name,
                            Team.name == team_name))
  if not found:
    raise InvalidTeamException('Team \'%s\' is not a team in org \'%s\'' %
                               (team_name, org_name))

  team = found[0]
  if team.role.name == 'admin':
    admin_teams = list(__get_user_admin_teams(org_name, team_name,
                                              removed_by_username))

    if len(admin_teams) <= 1:
      # The team we are trying to remove is the only admin team for this user
      msg = ('Deleting team \'%s\' would remove all admin from user \'%s\'' %
             (team_name, removed_by_username))
      raise DataModelException(msg)

  team.delete_instance(recursive=True, delete_nullable=True)


def add_user_to_team(user, team):
  try:
    return TeamMember.create(user=user, team=team)
  except Exception:
    raise DataModelException('Unable to add user \'%s\' to team: \'%s\'' %
                             (user.username, team.name))


def remove_user_from_team(org_name, team_name, username, removed_by_username):
  Org = User.alias()
  joined = TeamMember.select().join(User).switch(TeamMember).join(Team)
  with_role = joined.join(TeamRole)
  with_org = with_role.switch(Team).join(Org,
                                         on=(Org.id == Team.organization))
  found = list(with_org.where(User.username == username,
                              Org.username == org_name,
                              Team.name == team_name))

  if not found:
    raise DataModelException('User %s does not belong to team %s' %
                             (username, team_name))

  if username == removed_by_username:
    admin_team_query = __get_user_admin_teams(org_name, team_name, username)
    admin_team_names = {team.name for team in admin_team_query}
    if team_name in admin_team_names and len(admin_team_names) <= 1:
      msg = 'User cannot remove themselves from their only admin team.'
      raise DataModelException(msg)

  user_in_team = found[0]
  user_in_team.delete_instance()


def get_team_org_role(team):
  return TeamRole.get(TeamRole.id == team.role.id)


def set_team_org_permission(team, team_role_name, set_by_username):
  if team.role.name == 'admin' and team_role_name != 'admin':
    # We need to make sure we're not removing the users only admin role
    user_admin_teams = __get_user_admin_teams(team.organization.username,
                                              team.name, set_by_username)
    admin_team_set = {admin_team.name for admin_team in user_admin_teams}
    if team.name in admin_team_set and len(admin_team_set) <= 1:
      msg = (('Cannot remove admin from team \'%s\' because calling user ' +
              'would no longer have admin on org \'%s\'') %
             (team.name, team.organization.username))
      raise DataModelException(msg)

  new_role = TeamRole.get(TeamRole.name == team_role_name)
  team.role = new_role
  team.save()
  return team


def create_federated_user(username, email, service_name, service_id):
  new_user = create_user(username, None, email)
  new_user.verified = True
  new_user.save()

  service = LoginService.get(LoginService.name == service_name)
  FederatedLogin.create(user=new_user, service=service,
                        service_ident=service_id)

  return new_user


def attach_federated_login(user, service_name, service_id):
  service = LoginService.get(LoginService.name == service_name)
  FederatedLogin.create(user=user, service=service, service_ident=service_id)
  return user


def verify_federated_login(service_name, service_id):
  selected = FederatedLogin.select(FederatedLogin, User)
  with_service = selected.join(LoginService)
  with_user = with_service.switch(FederatedLogin).join(User)
  found = with_user.where(FederatedLogin.service_ident == service_id,
                          LoginService.name == service_name)

  found_list = list(found)

  if found_list:
    return found_list[0].user

  return None


def list_federated_logins(user):
  selected = FederatedLogin.select(FederatedLogin.service_ident,
                                   LoginService.name)
  joined = selected.join(LoginService)
  return joined.where(LoginService.name != 'quayrobot',
                      FederatedLogin.user == user)


def create_confirm_email_code(user, new_email=None):
  if new_email:
    if not validate_email(new_email):
      raise InvalidEmailAddressException('Invalid email address: %s' %
                                         new_email)

  code = EmailConfirmation.create(user=user, email_confirm=True,
                                  new_email=new_email)
  return code


def confirm_user_email(code):
  try:
    code = EmailConfirmation.get(EmailConfirmation.code == code,
                                 EmailConfirmation.email_confirm == True)
  except EmailConfirmation.DoesNotExist:
    raise DataModelException('Invalid email confirmation code.')    

  user = code.user
  user.verified = True

  new_email = code.new_email
  if new_email:
    if find_user_by_email(new_email):
      raise DataModelException('E-mail address already used.')    
      
    user.email = new_email
  
  user.save()

  code.delete_instance()

  return user, new_email


def create_reset_password_email_code(email):
  try:
    user = User.get(User.email == email)
  except User.DoesNotExist:
    raise InvalidEmailAddressException('Email address was not found.');

  if user.organization:
    raise InvalidEmailAddressException('Organizations can not have passwords.')

  code = EmailConfirmation.create(user=user, pw_reset=True)
  return code


def validate_reset_code(code):
  try:
    code = EmailConfirmation.get(EmailConfirmation.code == code,
                                 EmailConfirmation.pw_reset == True)
  except EmailConfirmation.DoesNotExist:
    return None

  user = code.user
  code.delete_instance()

  return user


def find_user_by_email(email):
  try:
    return User.get(User.email == email)
  except User.DoesNotExist:
    return None


def get_user(username):
  try:
    return User.get(User.username == username, User.organization == False)
  except User.DoesNotExist:
    return None


def get_user_or_org_by_customer_id(customer_id):
  try:
    return User.get(User.stripe_id == customer_id)
  except User.DoesNotExist:
    return None

def get_matching_teams(team_prefix, organization):  
  query = Team.select().where(Team.name ** (team_prefix + '%'),
                              Team.organization == organization)
  return query.limit(10)


def get_matching_users(username_prefix, robot_namespace=None,
                       organization=None):
  direct_user_query = (User.username ** (username_prefix + '%') & 
                       (User.organization == False) & (User.robot == False))

  if robot_namespace:
    robot_prefix = format_robot_username(robot_namespace, username_prefix)
    direct_user_query = (direct_user_query |
                         (User.username ** (robot_prefix + '%') & 
                          (User.robot == True)))

  query = (User
    .select(User.username, fn.Sum(Team.id), User.robot)
    .group_by(User.username)
    .where(direct_user_query))

  if organization:
    query = (query
      .join(TeamMember, JOIN_LEFT_OUTER)
      .join(Team, JOIN_LEFT_OUTER, on=((Team.id == TeamMember.team) &
                                       (Team.organization == organization))))


  class MatchingUserResult(object):
    def __init__(self, *args):
      self.username = args[0]
      self.is_robot = args[2]
      if organization:
        self.is_org_member = (args[1] != None)
      else:
        self.is_org_member = None


  return (MatchingUserResult(*args) for args in query.tuples().limit(10))


def verify_user(username_or_email, password):
  try:
    fetched = User.get((User.username == username_or_email) |
                       (User.email == username_or_email))
  except User.DoesNotExist:
    return None

  if (fetched.password_hash and 
      bcrypt.hashpw(password, fetched.password_hash) ==
      fetched.password_hash):
    return fetched

  # We weren't able to authorize the user
  return None


def get_user_organizations(username):
  UserAlias = User.alias()
  all_teams = User.select().distinct().join(Team).join(TeamMember)
  with_user = all_teams.join(UserAlias, on=(UserAlias.id == TeamMember.user))
  return with_user.where(User.organization == True,
                         UserAlias.username == username)


def get_organization(name):
  try:
    return User.get(username=name, organization=True)
  except User.DoesNotExist:
    raise InvalidOrganizationException('Organization does not exist: %s' %
                                       name)


def get_organization_team(orgname, teamname):
  joined = Team.select().join(User)
  query = joined.where(Team.name == teamname, User.organization == True,
                       User.username == orgname).limit(1)
  result = list(query)
  if not result:
    raise InvalidTeamException('Team does not exist: %s/%s', orgname,
                               teamname)

  return result[0]


def get_organization_members_with_teams(organization, membername = None):
  joined = TeamMember.select().annotate(Team).annotate(User)
  query = joined.where(Team.organization == organization)
  if membername:
    query = query.where(User.username == membername)
  return query

def get_organization_team_members(teamid):
  joined = User.select().join(TeamMember).join(Team)
  query = joined.where(Team.id == teamid)
  return query


def get_organization_member_set(orgname):
  Org = User.alias()
  user_teams = User.select(User.username).join(TeamMember).join(Team)
  with_org = user_teams.join(Org, on=(Org.username == orgname))
  return {user.username for user in with_org}


def get_teams_within_org(organization):
  return Team.select().where(Team.organization == organization)


def get_user_teams_within_org(username, organization):
  joined = Team.select().join(TeamMember).join(User)
  return joined.where(Team.organization == organization,
                      User.username == username)


def get_visible_repository_count(username=None, include_public=True,
                                 sort=False, namespace=None):
  return get_visible_repository_internal(username=username,
                                         include_public=include_public,
                                         sort=sort, namespace=namespace,
                                         get_count=True)

def get_visible_repositories(username=None, include_public=True, page=None,
                             limit=None, sort=False, namespace=None):
  return get_visible_repository_internal(username=username,
                                         include_public=include_public,
                                         page=page, limit=limit, sort=sort,
                                         namespace=namespace, get_count=False)


def get_visible_repository_internal(username=None, include_public=True,
                                    limit=None, page=None, sort=False,
                                    namespace=None, get_count=False):
  if not username and not include_public:
    return []

  query = (Repository
             .select()  # Note: We need to leave this blank for the get_count case. Otherwise, MySQL/RDS complains.
             .distinct()
             .join(Visibility)
             .switch(Repository)
             .join(RepositoryPermission, JOIN_LEFT_OUTER))

  where_clause = None
  admin_query = None
  if username:
    UserThroughTeam = User.alias()
    Org = User.alias()
    AdminTeam = Team.alias()
    AdminTeamMember = TeamMember.alias()
    AdminUser = User.alias()

    query = (query
      .join(User, JOIN_LEFT_OUTER)
      .switch(RepositoryPermission)
      .join(Team, JOIN_LEFT_OUTER)
      .join(TeamMember, JOIN_LEFT_OUTER)
      .join(UserThroughTeam, JOIN_LEFT_OUTER, on=(UserThroughTeam.id ==
                                                  TeamMember.user))
      .switch(Repository)
      .join(Org, JOIN_LEFT_OUTER, on=(Org.username == Repository.namespace))
      .join(AdminTeam, JOIN_LEFT_OUTER, on=(Org.id ==
                                            AdminTeam.organization))
      .join(TeamRole, JOIN_LEFT_OUTER, on=(AdminTeam.role == TeamRole.id))
      .switch(AdminTeam)
      .join(AdminTeamMember, JOIN_LEFT_OUTER, on=(AdminTeam.id ==
                                                  AdminTeamMember.team))
      .join(AdminUser, JOIN_LEFT_OUTER, on=(AdminTeamMember.user ==
                                            AdminUser.id)))

    where_clause = ((User.username == username) |
                    (UserThroughTeam.username == username) |
                    ((AdminUser.username == username) &
                     (TeamRole.name == 'admin')))

    if namespace:
      where_clause = where_clause & (Repository.namespace == namespace)

  if include_public:
    new_clause = (Visibility.name == 'public')
    if where_clause:
      where_clause = where_clause | new_clause
    else:
      where_clause = new_clause

  if sort:
    query = query.order_by(Repository.description.desc())

  if page:
    query = query.paginate(page, limit)
  elif limit:
    query = query.limit(limit)

  where = query.where(where_clause)
  if get_count:
    return where.count()
  else:
    return where


def get_matching_repositories(repo_term, username=None):
  namespace_term = repo_term
  name_term = repo_term

  visible = get_visible_repositories(username)

  search_clauses = (Repository.name ** ('%' + name_term + '%') |
                    Repository.namespace ** ('%' + namespace_term + '%'))

  # Handle the case where the user has already entered a namespace path.
  if repo_term.find('/') > 0:
    parts = repo_term.split('/', 1)
    namespace_term = '/'.join(parts[:-1])
    name_term = parts[-1]

    search_clauses = (Repository.name ** ('%' + name_term + '%') &
                      Repository.namespace ** ('%' + namespace_term + '%'))

  final = visible.where(search_clauses).limit(10)
  return list(final)


def change_password(user, new_password):
  if not validate_password(new_password):
    raise InvalidPasswordException(INVALID_PASSWORD_MESSAGE)

  pw_hash = bcrypt.hashpw(new_password, bcrypt.gensalt())
  user.password_hash = pw_hash
  user.save()


def change_invoice_email(user, invoice_email):
  user.invoice_email = invoice_email
  user.save()


def update_email(user, new_email):
  user.email = new_email
  user.verified = False
  user.save()


def get_all_user_permissions(user):
  select = RepositoryPermission.select(RepositoryPermission, Role, Repository)
  with_role = select.join(Role)
  with_repo = with_role.switch(RepositoryPermission).join(Repository)
  through_user = with_repo.switch(RepositoryPermission).join(User,
                                                             JOIN_LEFT_OUTER)
  as_perm = through_user.switch(RepositoryPermission)
  through_team = as_perm.join(Team, JOIN_LEFT_OUTER).join(TeamMember,
                                                          JOIN_LEFT_OUTER)

  UserThroughTeam = User.alias()
  with_team_member = through_team.join(UserThroughTeam, JOIN_LEFT_OUTER,
                                       on=(UserThroughTeam.id ==
                                           TeamMember.user))

  return with_team_member.where((User.id == user) | 
                                (UserThroughTeam.id == user))


def delete_prototype_permission(org, uid):
  found = get_prototype_permission(org, uid)
  if not found:
    return None

  found.delete_instance()
  return found


def get_prototype_permission(org, uid):
  try:
    return PermissionPrototype.get(PermissionPrototype.org == org,
                                   PermissionPrototype.uuid == uid)
  except PermissionPrototype.DoesNotExist:
    return None


def get_prototype_permissions(org):
  ActivatingUser = User.alias()
  DelegateUser = User.alias()
  query = (PermissionPrototype
    .select()
    .where(PermissionPrototype.org == org)
    .join(ActivatingUser, JOIN_LEFT_OUTER,
          on=(ActivatingUser.id == PermissionPrototype.activating_user))
    .join(DelegateUser, JOIN_LEFT_OUTER,
          on=(DelegateUser.id == PermissionPrototype.delegate_user))
    .join(Team, JOIN_LEFT_OUTER,
          on=(Team.id == PermissionPrototype.delegate_team))
    .join(Role, JOIN_LEFT_OUTER, on=(Role.id == PermissionPrototype.role)))
  return query


def update_prototype_permission(org, uid, role_name):
  found = get_prototype_permission(org, uid)
  if not found:
    return None

  new_role = Role.get(Role.name == role_name)
  found.role = new_role
  found.save()
  return found


def add_prototype_permission(org, role_name, activating_user,
                             delegate_user=None, delegate_team=None):
  new_role = Role.get(Role.name == role_name)
  return PermissionPrototype.create(org=org, role=new_role,
                                    activating_user=activating_user,
      delegate_user=delegate_user, delegate_team=delegate_team)


def get_org_wide_permissions(user):
  Org = User.alias()
  team_with_role = Team.select(Team, Org, TeamRole).join(TeamRole)
  with_org = team_with_role.switch(Team).join(Org, on=(Team.organization ==
                                                       Org.id))
  with_user = with_org.switch(Team).join(TeamMember).join(User)
  return with_user.where(User.id == user, Org.organization == True)


def get_all_repo_teams(namespace_name, repository_name):
  select = RepositoryPermission.select(Team.name.alias('team_name'),
                                       Role.name, RepositoryPermission)
  with_team = select.join(Team)
  with_role = with_team.switch(RepositoryPermission).join(Role)
  with_repo = with_role.switch(RepositoryPermission).join(Repository)
  return with_repo.where(Repository.namespace == namespace_name,
                         Repository.name == repository_name)


def get_all_repo_users(namespace_name, repository_name):
  select = RepositoryPermission.select(User.username, User.robot, Role.name,
                                       RepositoryPermission)
  with_user = select.join(User)
  with_role = with_user.switch(RepositoryPermission).join(Role)
  with_repo = with_role.switch(RepositoryPermission).join(Repository)
  return with_repo.where(Repository.namespace == namespace_name,
                         Repository.name == repository_name)


def get_repository_for_resource(resource_key):
  joined = Repository.select().join(RepositoryBuild)
  query = joined.where(RepositoryBuild.resource_key == resource_key).limit(1)
  result = list(query)
  if not result:
    return None

  return result[0]


def get_repository(namespace_name, repository_name):
  try:
    return Repository.get(Repository.name == repository_name,
                          Repository.namespace == namespace_name)
  except Repository.DoesNotExist:
    return None


def get_repo_image(namespace_name, repository_name, image_id):
  joined = Image.select().join(Repository)
  query = joined.where(Repository.name == repository_name,
                       Repository.namespace == namespace_name,
                       Image.docker_image_id == image_id).limit(1)
  result = list(query)
  if not result:
    return None

  return result[0]


def repository_is_public(namespace_name, repository_name):
  joined = Repository.select().join(Visibility)
  query = joined.where(Repository.namespace == namespace_name,
                       Repository.name == repository_name,
                       Visibility.name == 'public')
  return len(list(query)) > 0


def set_repository_visibility(repo, visibility):
  visibility_obj = Visibility.get(name=visibility)
  if not visibility_obj:
    return

  repo.visibility = visibility_obj
  repo.save()


def __apply_default_permissions(repo, proto_query, name_property,
                                create_permission_func):
  final_protos = {}
  for proto in proto_query:
    applies_to = proto.delegate_team or proto.delegate_user
    name = getattr(applies_to, name_property)
    # We will skip the proto if it is pre-empted by a more important proto
    if name in final_protos and proto.activating_user is None:
      continue

    # By this point, it is either a user specific proto, or there is no
    # proto yet, so we can safely assume it applies
    final_protos[name] = (applies_to, proto.role)

  for delegate, role in final_protos.values():
    create_permission_func(delegate, repo, role)


def create_repository(namespace, name, creating_user, visibility='private'):
  private = Visibility.get(name=visibility)
  repo = Repository.create(namespace=namespace, name=name,
                           visibility=private)
  admin = Role.get(name='admin')

  if creating_user and not creating_user.organization:
    RepositoryPermission.create(user=creating_user, repository=repo,
                                role=admin)

    if creating_user.username != namespace:
      # Permission prototypes only work for orgs
      org = get_organization(namespace)
      user_clause = ((PermissionPrototype.activating_user == creating_user) |
                     (PermissionPrototype.activating_user >> None))

      team_protos = (PermissionPrototype
                      .select()
                      .where(PermissionPrototype.org == org, user_clause,
                             PermissionPrototype.delegate_user >> None))

      def create_team_permission(team, repo, role):
        RepositoryPermission.create(team=team, repository=repo, role=role)

      __apply_default_permissions(repo, team_protos, 'name',
                                  create_team_permission)

      user_protos = (PermissionPrototype
                      .select()
                      .where(PermissionPrototype.org == org, user_clause,
                             PermissionPrototype.delegate_team >> None))

      def create_user_permission(user, repo, role):
        # The creating user always gets admin anyway
        if user.username == creating_user.username:
          return

        RepositoryPermission.create(user=user, repository=repo, role=role)

      __apply_default_permissions(repo, user_protos, 'username',
                                  create_user_permission)

  return repo


def create_image(docker_image_id, repository):
  new_image = Image.create(docker_image_id=docker_image_id,
                           repository=repository)
  return new_image


def set_image_checksum(docker_image_id, repository, checksum):
  fetched = Image.get(Image.docker_image_id == docker_image_id,
                      Image.repository == repository)
  fetched.checksum = checksum
  fetched.save()
  return fetched


def set_image_size(docker_image_id, namespace_name, repository_name,
                   image_size):
  joined = Image.select().join(Repository)
  image_list = list(joined.where(Repository.name == repository_name,
                                 Repository.namespace == namespace_name,
                                 Image.docker_image_id == docker_image_id))

  if not image_list:
    raise DataModelException('No image with specified id and repository')

  fetched = image_list[0]
  fetched.image_size = image_size
  fetched.save()
  return fetched


def set_image_metadata(docker_image_id, namespace_name, repository_name,
                       created_date_str, comment, command, parent=None):
  joined = Image.select().join(Repository)
  image_list = list(joined.where(Repository.name == repository_name,
                                 Repository.namespace == namespace_name,
                                 Image.docker_image_id == docker_image_id))

  if not image_list:
    raise DataModelException('No image with specified id and repository')

  fetched = image_list[0]
  fetched.created = dateutil.parser.parse(created_date_str)
  fetched.comment = comment
  fetched.command = command

  if parent:
    fetched.ancestors = '%s%s/' % (parent.ancestors, parent.id)

  fetched.save()
  return fetched


def get_repository_images(namespace_name, repository_name):
  joined = Image.select().join(Repository)
  return joined.where(Repository.name == repository_name,
                      Repository.namespace == namespace_name)


def list_repository_tags(namespace_name, repository_name):
  select = RepositoryTag.select(RepositoryTag, Image)
  with_repo = select.join(Repository)
  with_image = with_repo.switch(RepositoryTag).join(Image)
  return with_image.where(Repository.name == repository_name,
                          Repository.namespace == namespace_name)

def delete_tag_and_images(namespace_name, repository_name, tag_name):
  all_images = get_repository_images(namespace_name, repository_name)
  all_tags = list_repository_tags(namespace_name, repository_name)

  # Find the tag's information.
  found_tag = None
  for tag in all_tags:
    if tag.name == tag_name:
      found_tag = tag
      break

  if not found_tag:
    return

  # Build the set of database IDs corresponding to the tag's ancestor images,
  # as well as the tag's image itself.
  tag_image_ids = set(found_tag.image.ancestors.split('/'))
  tag_image_ids.add(str(found_tag.image.id))

  # Filter out any images that belong to any other tags.
  for tag in all_tags:
    if tag.name != tag_name:
      # Remove all ancestors of the tag.
      tag_image_ids = tag_image_ids - set(tag.image.ancestors.split('/'))

      # Remove the current image ID.
      tag_image_ids.discard(str(tag.image.id))

  # Find all the images that belong to the tag.
  tag_images = [image for image in all_images
                if str(image.id) in tag_image_ids]

  # Delete the tag found.
  found_tag.delete_instance()

  # Delete the images found.
  for image in tag_images:
    image.delete_instance()

    repository_path = store.image_path(namespace_name, repository_name,
                                       image.docker_image_id)
    logger.debug('Recursively deleting image path: %s' % repository_path)
    store.remove(repository_path)


def garbage_collect_repository(namespace_name, repository_name):
  # Get a list of all images used by tags in the repository
  tag_query = (RepositoryTag
    .select(RepositoryTag, Image)
    .join(Image)
    .switch(RepositoryTag)
    .join(Repository)
    .where(Repository.name == repository_name,
           Repository.namespace == namespace_name))

  referenced_anscestors = set()
  for tag in tag_query:
    # The anscestor list is in the format '/1/2/3/', extract just the ids
    anscestor_id_strings = tag.image.ancestors.split('/')[1:-1]
    ancestor_list = [int(img_id_str) for img_id_str in anscestor_id_strings]
    referenced_anscestors = referenced_anscestors.union(set(ancestor_list))
    referenced_anscestors.add(tag.image.id)

  all_repo_images = get_repository_images(namespace_name, repository_name)
  all_images = {int(img.id):img for img in all_repo_images}
  to_remove = set(all_images.keys()).difference(referenced_anscestors)

  logger.info('Cleaning up unreferenced images: %s', to_remove)

  for image_id_to_remove in to_remove:
    image_to_remove = all_images[image_id_to_remove]
    image_path = store.image_path(namespace_name, repository_name,
                                       image_to_remove.docker_image_id)
    image_to_remove.delete_instance()
    logger.debug('Deleting image storage: %s' % image_path)
    store.remove(image_path)

  return len(to_remove)


def get_tag_image(namespace_name, repository_name, tag_name):
  joined = Image.select().join(RepositoryTag).join(Repository)
  fetched = list(joined.where(Repository.name == repository_name,
                              Repository.namespace == namespace_name,
                              RepositoryTag.name == tag_name))

  if not fetched:
    raise DataModelException('Unable to find image for tag.')

  return fetched[0]


def get_image_by_id(namespace_name, repository_name, docker_image_id):
  joined = Image.select().join(Repository)
  fetched = list(joined.where(Repository.name == repository_name,
                              Repository.namespace == namespace_name,
                              Image.docker_image_id == docker_image_id))

  if not fetched:
    raise DataModelException('Unable to find image \'%s\' for repo \'%s/%s\'' %
                             (docker_image_id, namespace_name,
                              repository_name))

  return fetched[0]


def get_parent_images(image_obj):
  """ Returns a list of parent Image objects in chronilogical order. """
  parents = image_obj.ancestors
  parent_db_ids = parents.strip('/').split('/')

  if parent_db_ids == ['']:
    return []

  or_clauses = [(Image.id == db_id) for db_id in parent_db_ids]
  parent_images = Image.select().where(reduce(operator.or_, or_clauses))
  id_to_image = {unicode(image.id): image for image in parent_images}

  return [id_to_image[parent_id] for parent_id in parent_db_ids]


def create_or_update_tag(namespace_name, repository_name, tag_name,
                         tag_docker_image_id):
  try:
    repo = Repository.get(Repository.name == repository_name,
                          Repository.namespace == namespace_name)
  except Repository.DoesNotExist:
    raise DataModelException('Invalid repository %s/%s' %
                             (namespace_name, repository_name))

  try:
    image = Image.get(Image.docker_image_id == tag_docker_image_id,
                      Image.repository == repo)
  except Image.DoesNotExist:
    raise DataModelException('Invalid image with id: %s' %
                             tag_docker_image_id)

  try:
    tag = RepositoryTag.get(RepositoryTag.repository == repo,
                            RepositoryTag.name == tag_name)
    tag.image = image
    tag.save()
  except RepositoryTag.DoesNotExist:
    tag = RepositoryTag.create(repository=repo, image=image, name=tag_name)

  return tag


def delete_tag(namespace_name, repository_name, tag_name):
  joined = RepositoryTag.select().join(Repository)
  found = list(joined.where(Repository.name == repository_name,
                            Repository.namespace == namespace_name,
                            RepositoryTag.name == tag_name))

  if not found:
    msg = ('Invalid repository tag \'%s\' on repository \'%s/%s\'' %
           (tag_name, namespace_name, repository_name))
    raise DataModelException(msg)

  found[0].delete_instance()


def delete_all_repository_tags(namespace_name, repository_name):
  try:
    repo = Repository.get(Repository.name == repository_name,
                          Repository.namespace == namespace_name)
  except Repository.DoesNotExist:
    raise DataModelException('Invalid repository \'%s/%s\'' %
                             (namespace_name, repository_name))
  RepositoryTag.delete().where(RepositoryTag.repository == repo).execute()


def __entity_permission_repo_query(entity_id, entity_table,
                                   entity_id_property, namespace_name,
                                   repository_name):
  """ This method works for both users and teams. """
  selected = RepositoryPermission.select(entity_table, Repository, Role,
                                         RepositoryPermission)
  with_user = selected.join(entity_table)
  with_role = with_user.switch(RepositoryPermission).join(Role)
  with_repo = with_role.switch(RepositoryPermission).join(Repository)
  return with_repo.where(Repository.name == repository_name,
                         Repository.namespace == namespace_name,
                         entity_id_property == entity_id)


def get_user_reponame_permission(username, namespace_name, repository_name):
  fetched = list(__entity_permission_repo_query(username, User, User.username,
                                                namespace_name,
                                                repository_name))
  if not fetched:
    raise DataModelException('User does not have permission for repo.')

  return fetched[0]


def get_team_reponame_permission(team_name, namespace_name, repository_name):
  fetched = list(__entity_permission_repo_query(team_name, Team, Team.name,
                                                namespace_name,
                                                repository_name))
  if not fetched:
    raise DataModelException('Team does not have permission for repo.')

  return fetched[0]


def delete_user_permission(username, namespace_name, repository_name):
  if username == namespace_name:
    raise DataModelException('Namespace owner must always be admin.')

  fetched = list(__entity_permission_repo_query(username, User, User.username,
                                                namespace_name,
                                                repository_name))
  if not fetched:
    raise DataModelException('User does not have permission for repo.')

  fetched[0].delete_instance()


def delete_team_permission(team_name, namespace_name, repository_name):
  fetched = list(__entity_permission_repo_query(team_name, Team, Team.name,
                                                namespace_name,
                                                repository_name))
  if not fetched:
    raise DataModelException('Team does not have permission for repo.')

  fetched[0].delete_instance()


def __set_entity_repo_permission(entity, permission_entity_property,
                                 namespace_name, repository_name, role_name):
  repo = Repository.get(Repository.name == repository_name,
                        Repository.namespace == namespace_name)
  new_role = Role.get(Role.name == role_name)

  # Fetch any existing permission for this entity on the repo
  try:
    entity_attr = getattr(RepositoryPermission, permission_entity_property)
    perm = RepositoryPermission.get(entity_attr == entity,
                                    RepositoryPermission.repository == repo)
    perm.role = new_role
    perm.save()
    return perm
  except RepositoryPermission.DoesNotExist:
    set_entity_kwargs = {permission_entity_property: entity}
    new_perm = RepositoryPermission.create(repository=repo, role=new_role,
                                           **set_entity_kwargs)
    return new_perm


def set_user_repo_permission(username, namespace_name, repository_name,
                             role_name):
  if username == namespace_name:
    raise DataModelException('Namespace owner must always be admin.')

  user = User.get(User.username == username)
  return __set_entity_repo_permission(user, 'user', namespace_name,
                                      repository_name, role_name)


def set_team_repo_permission(team_name, namespace_name, repository_name,
                             role_name):
  team = list(Team.select().join(User).where(Team.name == team_name,
                                             User.username == namespace_name))
  if not team:
    raise DataModelException('No team \'%s\' in organization \'%s\'.' %
                             (team_name, namespace_name))

  return __set_entity_repo_permission(team[0], 'team', namespace_name,
                                      repository_name, role_name)


def purge_repository(namespace_name, repository_name):
  fetched = Repository.get(Repository.name == repository_name,
                           Repository.namespace == namespace_name)
  fetched.delete_instance(recursive=True)

  repository_path = store.repository_namespace_path(namespace_name,
                                                    repository_name)
  logger.debug('Recursively deleting path: %s' % repository_path)
  store.remove(repository_path)


def get_private_repo_count(username):
  joined = Repository.select().join(Visibility)
  return joined.where(Repository.namespace == username,
                      Visibility.name == 'private').count()


def create_access_token(repository, role):
  role = Role.get(Role.name == role)
  new_token = AccessToken.create(repository=repository, temporary=True,
                                 role=role)
  return new_token


def create_delegate_token(namespace_name, repository_name, friendly_name):
  read_only = Role.get(name='read')
  repo = Repository.get(Repository.name == repository_name,
                        Repository.namespace == namespace_name)
  new_token = AccessToken.create(repository=repo, role=read_only,
                                 friendly_name=friendly_name, temporary=False)
  return new_token


def get_repository_delegate_tokens(namespace_name, repository_name):
  selected = AccessToken.select(AccessToken, Role)
  with_repo = selected.join(Repository)
  with_role = with_repo.switch(AccessToken).join(Role)
  return with_role.where(Repository.name == repository_name,
                         Repository.namespace == namespace_name,
                         AccessToken.temporary == False)


def get_repo_delegate_token(namespace_name, repository_name, code):
  repo_query = get_repository_delegate_tokens(namespace_name, repository_name)
  found = list(repo_query.where(AccessToken.code == code))

  if found:
    return found[0]
  else:
    raise InvalidTokenException('Unable to find token with code: %s' % code)


def set_repo_delegate_token_role(namespace_name, repository_name, code, role):
  token = get_repo_delegate_token(namespace_name, repository_name, code)

  if role != 'read' and role != 'write':
    raise DataModelException('Invalid role for delegate token: %s' % role)

  new_role = Role.get(Role.name == role)
  token.role = new_role
  token.save()

  return token


def delete_delegate_token(namespace_name, repository_name, code):
  token = get_repo_delegate_token(namespace_name, repository_name, code)
  token.delete_instance()
  return token


def load_token_data(code):
  """ Load the permissions for any token by code. """
  selected = AccessToken.select(AccessToken, Repository, Role)
  with_role = selected.join(Role)
  with_repo = with_role.switch(AccessToken).join(Repository)
  fetched = list(with_repo.where(AccessToken.code == code))

  if fetched:
    return fetched[0]
  else:
    raise InvalidTokenException('Invalid delegate token code: %s' % code)


def get_repository_build(namespace_name, repository_name, build_uuid):
  try:
    query = list_repository_builds(namespace_name, repository_name)
    return query.where(RepositoryBuild.uuid == build_uuid).get()

  except RepositoryBuild.DoesNotExist:
    msg = 'Unable to locate a build by id: %s' % build_uuid
    raise InvalidRepositoryBuildException(msg)


def list_repository_builds(namespace_name, repository_name,
                           include_inactive=True):
  query = (RepositoryBuild
    .select(RepositoryBuild, RepositoryBuildTrigger, BuildTriggerService)
    .join(Repository)
    .switch(RepositoryBuild)
    .join(RepositoryBuildTrigger, JOIN_LEFT_OUTER)
    .join(BuildTriggerService, JOIN_LEFT_OUTER)
    .where(Repository.name == repository_name,
           Repository.namespace == namespace_name))

  if not include_inactive:
    query = query.where(RepositoryBuild.phase != 'error',
                        RepositoryBuild.phase != 'complete')

  return query


def create_repository_build(repo, access_token, resource_key, tag,
                            display_name, trigger=None):
  return RepositoryBuild.create(repository=repo, access_token=access_token,
                                resource_key=resource_key, tag=tag,
                                display_name=display_name, trigger=trigger)


def create_webhook(repo, params_obj):
  return Webhook.create(repository=repo, parameters=json.dumps(params_obj))


def get_webhook(namespace_name, repository_name, public_id):
  joined = Webhook.select().join(Repository)
  found = list(joined.where(Repository.namespace == namespace_name,
                            Repository.name == repository_name,
                            Webhook.public_id == public_id))

  if not found:
    raise InvalidWebhookException('No webhook found with id: %s' % public_id)

  return found[0]


def list_webhooks(namespace_name, repository_name):
  joined = Webhook.select().join(Repository)
  return joined.where(Repository.namespace == namespace_name,
                      Repository.name == repository_name)


def delete_webhook(namespace_name, repository_name, public_id):
  webhook = get_webhook(namespace_name, repository_name, public_id)
  webhook.delete_instance()
  return webhook


def list_logs(user_or_organization_name, start_time, end_time, performer=None,
              repository=None):
  joined = LogEntry.select().join(User)
  if repository:
    joined = joined.where(LogEntry.repository == repository)

  if performer:
    joined = joined.where(LogEntry.performer == performer)

  return joined.where(
    User.username == user_or_organization_name,
    LogEntry.datetime >= start_time,
    LogEntry.datetime < end_time).order_by(LogEntry.datetime.desc())


def log_action(kind_name, user_or_organization_name, performer=None,
               repository=None, access_token=None, ip=None, metadata={},
               timestamp=None):
  if not timestamp:
    timestamp = datetime.today()

  kind = LogEntryKind.get(LogEntryKind.name == kind_name)
  account = User.get(User.username == user_or_organization_name)
  LogEntry.create(kind=kind, account=account, performer=performer,
                  repository=repository, access_token=access_token, ip=ip,
                  metadata_json=json.dumps(metadata), datetime=timestamp)


def create_build_trigger(repo, service_name, auth_token, user):
  service = BuildTriggerService.get(name=service_name)
  trigger = RepositoryBuildTrigger.create(repository=repo, service=service,
                                          auth_token=auth_token,
                                          connected_user=user)
  return trigger


def get_build_trigger(namespace_name, repository_name, trigger_uuid):
  try:
    return (RepositoryBuildTrigger
      .select(RepositoryBuildTrigger, BuildTriggerService, Repository)
      .join(BuildTriggerService)
      .switch(RepositoryBuildTrigger)
      .join(Repository)
      .switch(RepositoryBuildTrigger)
      .join(User)
      .where(RepositoryBuildTrigger.uuid == trigger_uuid,
             Repository.namespace == namespace_name,
             Repository.name == repository_name)
      .get())
  except RepositoryBuildTrigger.DoesNotExist:
    msg = 'No build trigger with uuid: %s' % trigger_uuid
    raise InvalidBuildTriggerException(msg)


def list_build_triggers(namespace_name, repository_name):
  return (RepositoryBuildTrigger
    .select(RepositoryBuildTrigger, BuildTriggerService, Repository)
    .join(BuildTriggerService)
    .switch(RepositoryBuildTrigger)
    .join(Repository)
    .where(Repository.namespace == namespace_name,
           Repository.name == repository_name))


def delete_build_trigger(namespace_name, repository_name, trigger_uuid):
  trigger = get_build_trigger(namespace_name, repository_name, trigger_uuid)
  trigger.delete_instance()


def list_trigger_builds(namespace_name, repository_name, trigger_uuid,
                        limit=None):
  query = (list_repository_builds(namespace_name, repository_name)
    .where(RepositoryBuildTrigger.uuid == trigger_uuid))

  if limit:
    return query.limit(limit)
  else:
    return query