1 files changed, 30 insertions, 34 deletions

diff --git a/coip/apps/membership/models.py b/coip/apps/membership/models.py
index b19c04a..71143b5 100644
--- a/coip/apps/membership/models.py
+++ b/coip/apps/membership/models.py
@@ -7,18 +7,23 @@ from django.db import models
 from django.contrib.auth.models import User
 from coip.apps.name.models import Name
 import datetime
-import logging
-from coip.apps.entity.models import Entity
 import tagging
 from django.core.mail import send_mail
 from coip.settings import NOREPLY
+from coip.extensions.templatetags.userdisplay import userdisplay
+from coip.apps.userprofile.models import UserProfile
+
+STATUS = {UserProfile.INTERNAL:'internal',
+          UserProfile.ENTITY:'entity',
+          UserProfile.SSHKEY:'sshkey',
+          UserProfile.X509:'certificate',
+          UserProfile.FEDID:'fedid'}
 
 class Membership(models.Model):
     '''
     Membership in a namespace/group
     '''
-    user = models.ForeignKey(User,blank=True,null=True,related_name='memberships')
-    entity = models.ForeignKey(Entity,blank=True,null=True,related_name='entity') # XXX should this be memberships too?
+    user = models.ForeignKey(User,related_name='memberships')
     name = models.ForeignKey(Name,related_name='memberships')
     enabled = models.BooleanField()
     hidden = models.BooleanField()
@@ -27,12 +32,7 @@ class Membership(models.Model):
     expires = models.DateTimeField(blank=True,null=True)
     
     def __unicode__(self):
-        who = "(unknown)"
-        if self.user:
-            who = self.user
-        elif self.entity:
-            who = self.entity
-            
+        who = self.user    
         status = ""
         if not self.enabled:
             status = " (disabled)"
@@ -49,18 +49,15 @@ class Membership(models.Model):
             return "active"
         else:
             return "inactive";
-        
-    def is_user(self):
-        return self.user != None
     
-    def is_entity(self):
-        return self.entity != None
+    def type(self):
+        return STATUS[self.user.get_profile().type]
     
     def send_notification(self,what):
         if not self.user or not self.user.email:
             return
-                
-        send_mail('You have been %s \'%s\'' % (what,self.name.short),
+        
+        send_mail('%s have been %s \'%s\'' % (userdisplay(self.user),what,self.name.short),
                   '''
 You have been %s \'%s\'.
 
@@ -74,41 +71,40 @@ To view information about \'%s\' open this link in your browser:
         return
     
 def add_member(name,member_name,hidden=False):
-    if isinstance(member_name,User):
-        (m,created)  = Membership.objects.get_or_create(user=member_name,name=name)
-    else:
-        (m,created)  = Membership.objects.get_or_create(entity=member_name,name=name)
-        
+    (m,created)  = Membership.objects.get_or_create(user=member_name,name=name)
     if created or not m.enabled or m.hidden != hidden:
         m.enabled = True
         m.hidden = hidden
         m.save()
+        
+    if name.nmembers != -1:
+        name.nmembers = -1
+        name.save()
     
     return m.send_notification("added to")
         
 def disable_member(name,member_name):
-    if isinstance(member_name,User):
-        m = Membership.objects.get(name=name,user=member_name)
-    else:
-        m = Membership.objects.get(name=name,entity=member_name)
+    m = Membership.objects.get(name=name,user=member_name)
     if m:
         m.enabled = False
         m.save()
         m.send_notification("temporarily removed from")
         
+    if name.nmembers != -1:
+        name.nmembers = -1
+        name.save()
+        
 def remove_member(name,member_name):
-    if isinstance(member_name,User):
-        m = Membership.objects.get(name=name,user=member_name)
-    else:
-        m = Membership.objects.get(name=name,entity=member_name)
+    m = Membership.objects.get(name=name,user=member_name)
     if m:
         m.send_notification("removed from")
         m.delete()
+    
+    if name.nmembers != -1:
+        name.nmembers = -1
+        name.save()
         
 def has_member(name,member_name):
-    if isinstance(member_name,User):
-        return Membership.objects.filter(name=name,user=member_name)
-    else:
-        return Membership.objects.filter(name=name,entity=member_name)
+    return Membership.objects.filter(name=name,user=member_name)
 
 tagging.register(Membership)
\ No newline at end of file