From c9e192fae46f46a80baf64b4ff478985c39e34f0 Mon Sep 17 00:00:00 2001
From: liuzheng712 <liuzheng712@gmail.com>
Date: Wed, 7 Oct 2015 12:17:20 +0800
Subject: [PATCH] recovery User's function

---
 juser/models.py | 84 +++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 81 insertions(+), 3 deletions(-)

diff --git a/juser/models.py b/juser/models.py
index 41dd705a..2d38e341 100644
--- a/juser/models.py
+++ b/juser/models.py
@@ -1,6 +1,7 @@
 # coding: utf-8
 
 from django.db import models
+from django.contrib.auth.models import AbstractUser
 
 
 class UserGroup(models.Model):
@@ -19,9 +20,6 @@ class UserGroup(models.Model):
             self.save()
 
 
-from django.contrib.auth.models import AbstractUser
-
-
 class User(AbstractUser):
     USER_ROLE_CHOICES = (
         ('SU', 'SuperUser'),
@@ -34,6 +32,86 @@ class User(AbstractUser):
     group = models.ManyToManyField(UserGroup)
     ssh_key_pwd = models.CharField(max_length=200)
 
+    def __unicode__(self):
+        return self.username
+
+    def get_asset_group(self):
+        """
+        Get user host_groups.
+        获取用户有权限的主机组
+        """
+        host_group_list = []
+        perm_list = []
+        user_group_all = self.group.all()
+        for user_group in user_group_all:
+            perm_list.extend(user_group.perm_set.all())
+        for perm in perm_list:
+            host_group_list.append(perm.asset_group)
+        return host_group_list
+
+    def get_asset_group_info(self, printable=False):
+        """
+        Get or print asset group info
+        获取或打印用户授权资产组
+        """
+        asset_groups_info = {}
+        asset_groups = self.get_asset_group()
+        for asset_group in asset_groups:
+            asset_groups_info[asset_group.id] = [asset_group.name, asset_group.comment]
+        if printable:
+            for group_id in asset_groups_info:
+                if asset_groups_info[group_id][1]:
+                    print "[%3s] %s -- %s" % (group_id,
+                                              asset_groups_info[group_id][0],
+                                              asset_groups_info[group_id][1])
+                else:
+                    print "[%3s] %s" % (group_id, asset_groups_info[group_id][0])
+                print ''
+        else:
+            return asset_groups_info
+
+    def get_asset(self):
+        """
+        Get the assets of under the user control.
+        获取主机列表
+        """
+        assets = []
+        asset_groups = self.get_asset_group()
+        for asset_group in asset_groups:
+            assets.extend(asset_group.asset_set.all())
+        return assets
+
+    def get_asset_info(self, printable=False):
+        """
+        Get or print the user asset info
+        获取或打印用户资产信息
+        """
+        from jasset.models import AssetAlias
+        assets_info = {}
+        assets = self.get_asset()
+        for asset in assets:
+            asset_alias = AssetAlias.objects.filter(user=self, asset=asset)
+            if asset_alias and asset_alias[0].alias != '':
+                assets_info[asset.ip] = [asset.id, asset.ip, str(asset_alias[0].alias)]
+            else:
+                assets_info[asset.ip] = [asset.id, asset.ip, str(asset.comment)]
+        if printable:
+            ips = assets_info.keys()
+            ips.sort()
+            for ip in ips:
+                if assets_info[ip][2]:
+                    print '%-15s -- %s' % (ip, assets_info[ip][2])
+                else:
+                    print '%-15s' % ip
+            print ''
+        else:
+            return assets_info
+
+    def update(self, **kwargs):
+        for key, value in kwargs.items():
+            self.__setattr__(key, value)
+            self.save()
+
 
 class AdminGroup(models.Model):
     """
-- 
2.18.0