加入权限管理

backend/constants/auth.go

@@ -0,0 +1,7 @@
+package constants
+
+const (
+	OwnerTypeAll    = "all"
+	OwnerTypeMe     = "me"
+	OwnerTypePublic = "public"
+)

backend/main.go

@@ -100,6 +100,14 @@ func main() {
 			panic(err)
 		}
 		log.Info("initialized challenge service successfully")
+
+		// 初始化清理服务
+		if err := services.InitCleanService(); err != nil {
+			log.Error("init clean service error:" + err.Error())
+			debug.PrintStack()
+			panic(err)
+		}
+		log.Info("initialized clean service successfully")
 	}
 
 	// 初始化任务执行器

backend/model/project.go

@@ -16,11 +16,12 @@ type Project struct {
 	Tags        []string      `json:"tags" bson:"tags"`
 
 	// 前端展示
-	Spiders []Spider `json:"spiders" bson:"spiders"`
+	Spiders  []Spider `json:"spiders" bson:"spiders"`
+	Username string   `json:"username" bson:"username"`
 
 	UserId   bson.ObjectId `json:"user_id" bson:"user_id"`
-	CreateTs time.Time `json:"create_ts" bson:"create_ts"`
-	UpdateTs time.Time `json:"update_ts" bson:"update_ts"`
+	CreateTs time.Time     `json:"create_ts" bson:"create_ts"`
+	UpdateTs time.Time     `json:"update_ts" bson:"update_ts"`
 }
 
 func (p *Project) Save() error {
@@ -90,15 +91,21 @@ func GetProject(id bson.ObjectId) (Project, error) {
 	return p, nil
 }
 
-func GetProjectList(filter interface{}, skip int, sortKey string) ([]Project, error) {
+func GetProjectList(filter interface{}, sortKey string) ([]Project, error) {
 	s, c := database.GetCol("projects")
 	defer s.Close()
 
 	var projects []Project
-	if err := c.Find(filter).Skip(skip).Limit(constants.Infinite).Sort(sortKey).All(&projects); err != nil {
+	if err := c.Find(filter).Sort(sortKey).All(&projects); err != nil {
 		debug.PrintStack()
 		return projects, err
 	}
+
+	for i, p := range projects {
+		// 获取用户名称
+		user, _ := GetUser(p.UserId)
+		projects[i].Username = user.Username
+	}
 	return projects, nil
 }
 

backend/model/schedule.go

@@ -29,6 +29,7 @@ type Schedule struct {
 
 	// 前端展示
 	SpiderName string `json:"spider_name" bson:"spider_name"`
+	Username   string `json:"user_name" bson:"user_name"`
 	Nodes      []Node `json:"nodes" bson:"nodes"`
 	Message    string `json:"message" bson:"message"`
 
@@ -83,6 +84,10 @@ func GetScheduleList(filter interface{}) ([]Schedule, error) {
 			schedule.SpiderName = spider.Name
 		}
 
+		// 获取用户名称
+		user, _ := GetUser(schedule.UserId)
+		schedule.Username = user.Username
+
 		schs = append(schs, schedule)
 	}
 	return schs, nil
@@ -92,11 +97,16 @@ func GetSchedule(id bson.ObjectId) (Schedule, error) {
 	s, c := database.GetCol("schedules")
 	defer s.Close()
 
-	var result Schedule
-	if err := c.FindId(id).One(&result); err != nil {
-		return result, err
+	var schedule Schedule
+	if err := c.FindId(id).One(&schedule); err != nil {
+		return schedule, err
 	}
-	return result, nil
+
+	// 获取用户名称
+	user, _ := GetUser(schedule.UserId)
+	schedule.Username = user.Username
+
+	return schedule, nil
 }
 
 func UpdateSchedule(id bson.ObjectId, item Schedule) error {
@@ -147,11 +157,11 @@ func RemoveSchedule(id bson.ObjectId) error {
 	return nil
 }
 
-func GetScheduleCount() (int, error) {
+func GetScheduleCount(filter interface{}) (int, error) {
 	s, c := database.GetCol("schedules")
 	defer s.Close()
 
-	count, err := c.Count()
+	count, err := c.Find(filter).Count()
 	if err != nil {
 		return 0, err
 	}

backend/model/spider.go

@@ -33,6 +33,7 @@ type Spider struct {
 	Remark      string        `json:"remark" bson:"remark"`             // 备注
 	Src         string        `json:"src" bson:"src"`                   // 源码位置
 	ProjectId   bson.ObjectId `json:"project_id" bson:"project_id"`     // 项目ID
+	IsPublic    bool          `json:"is_public" bson:"is_public"`       // 是否公开
 
 	// 自定义爬虫
 	Cmd string `json:"cmd" bson:"cmd"` // 执行命令
@@ -63,6 +64,7 @@ type Spider struct {
 	LastStatus  string                  `json:"last_status"`  // 最后执行状态
 	Config      entity.ConfigSpiderData `json:"config"`       // 可配置爬虫配置
 	LatestTasks []Task                  `json:"latest_tasks"` // 最近任务列表
+	Username    string                  `json:"username""`
 
 	// 时间
 	UserId   bson.ObjectId `json:"user_id" bson:"user_id"`
@@ -83,6 +85,7 @@ func (spider *Spider) Save() error {
 	}
 
 	if err := c.UpdateId(spider.Id, spider); err != nil {
+		log.Errorf(err.Error())
 		debug.PrintStack()
 		return err
 	}
@@ -182,10 +185,22 @@ func GetSpiderList(filter interface{}, skip int, limit int, sortStr string) ([]S
 			continue
 		}
 
+		// 获取用户
+		var user User
+		if spider.UserId.Valid() {
+			user, err = GetUser(spider.UserId)
+			if err != nil {
+				log.Errorf(err.Error())
+				debug.PrintStack()
+				continue
+			}
+		}
+
 		// 赋值
 		spiders[i].LastRunTs = task.CreateTs
 		spiders[i].LastStatus = task.Status
 		spiders[i].LatestTasks = latestTasks
+		spiders[i].Username = user.Username
 	}
 
 	count, _ := c.Find(filter).Count()
@@ -221,13 +236,21 @@ func GetSpiderByName(name string) Spider {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
 
-	var result Spider
-	if err := c.Find(bson.M{"name": name}).One(&result); err != nil && err != mgo.ErrNotFound {
+	var spider Spider
+	if err := c.Find(bson.M{"name": name}).One(&spider); err != nil && err != mgo.ErrNotFound {
 		log.Errorf("get spider error: %s, spider_name: %s", err.Error(), name)
 		//debug.PrintStack()
-		return result
+		return spider
 	}
-	return result
+
+	// 获取用户
+	var user User
+	if spider.UserId.Valid() {
+		user, _ = GetUser(spider.UserId)
+	}
+	spider.Username = user.Username
+
+	return spider
 }
 
 // 获取爬虫(根据ID)
@@ -253,6 +276,14 @@ func GetSpider(id bson.ObjectId) (Spider, error) {
 		}
 		spider.Config = config
 	}
+
+	// 获取用户名称
+	var user User
+	if spider.UserId.Valid() {
+		user, _ = GetUser(spider.UserId)
+	}
+	spider.Username = user.Username
+
 	return spider, nil
 }
 
@@ -324,11 +355,11 @@ func RemoveAllSpider() error {
 }
 
 // 获取爬虫总数
-func GetSpiderCount() (int, error) {
+func GetSpiderCount(filter interface{}) (int, error) {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
 
-	count, err := c.Count()
+	count, err := c.Find(filter).Count()
 	if err != nil {
 		return 0, err
 	}

backend/model/task.go

@@ -31,6 +31,7 @@ type Task struct {
 	// 前端数据
 	SpiderName string `json:"spider_name"`
 	NodeName   string `json:"node_name"`
+	Username   string `json:"username"`
 
 	UserId   bson.ObjectId `json:"user_id" bson:"user_id"`
 	CreateTs time.Time     `json:"create_ts" bson:"create_ts"`
@@ -128,6 +129,10 @@ func GetTaskList(filter interface{}, skip int, limit int, sortKey string) ([]Tas
 		if node, err := task.GetNode(); err == nil {
 			tasks[i].NodeName = node.Name
 		}
+
+		// 获取用户名称
+		user, _ := GetUser(task.UserId)
+		task.Username = user.Username
 	}
 	return tasks, nil
 }
@@ -156,6 +161,11 @@ func GetTask(id string) (Task, error) {
 		debug.PrintStack()
 		return task, err
 	}
+
+	// 获取用户名称
+	user, _ := GetUser(task.UserId)
+	task.Username = user.Username
+
 	return task, nil
 }
 

backend/routes/project.go

@@ -4,6 +4,7 @@ import (
 	"crawlab/constants"
 	"crawlab/database"
 	"crawlab/model"
+	"crawlab/services"
 	"github.com/gin-gonic/gin"
 	"github.com/globalsign/mgo/bson"
 	"net/http"
@@ -18,8 +19,11 @@ func GetProjectList(c *gin.Context) {
 		query["tags"] = tag
 	}
 
+	// 获取校验
+	query = services.GetAuthQuery(query, c)
+
 	// 获取列表
-	projects, err := model.GetProjectList(query, 0, "+_id")
+	projects, err := model.GetProjectList(query, "+_id")
 	if err != nil {
 		HandleError(http.StatusInternalServerError, c, err)
 		return
@@ -74,6 +78,9 @@ func PutProject(c *gin.Context) {
 		return
 	}
 
+	// UserId
+	p.UserId = services.GetCurrentUserId(c)
+
 	if err := p.Add(); err != nil {
 		HandleError(http.StatusInternalServerError, c, err)
 		return

backend/routes/schedule.go

@@ -9,7 +9,12 @@ import (
 )
 
 func GetScheduleList(c *gin.Context) {
-	results, err := model.GetScheduleList(nil)
+	query := bson.M{}
+
+	// 获取校验
+	query = services.GetAuthQuery(query, c)
+
+	results, err := model.GetScheduleList(query)
 	if err != nil {
 		HandleError(http.StatusInternalServerError, c, err)
 		return

backend/routes/spider.go

@@ -29,13 +29,14 @@ import (
 // ======== 爬虫管理 ========
 
 func GetSpiderList(c *gin.Context) {
-	pageNum, _ := c.GetQuery("page_num")
-	pageSize, _ := c.GetQuery("page_size")
-	keyword, _ := c.GetQuery("keyword")
-	pid, _ := c.GetQuery("project_id")
-	t, _ := c.GetQuery("type")
-	sortKey, _ := c.GetQuery("sort_key")
-	sortDirection, _ := c.GetQuery("sort_direction")
+	pageNum := c.Query("page_num")
+	pageSize := c.Query("page_size")
+	keyword := c.Query("keyword")
+	pid := c.Query("project_id")
+	t := c.Query("type")
+	sortKey := c.Query("sort_key")
+	sortDirection := c.Query("sort_direction")
+	ownerType := c.Query("owner_type")
 
 	// 筛选-名称
 	filter := bson.M{
@@ -65,6 +66,21 @@ func GetSpiderList(c *gin.Context) {
 		filter["project_id"] = bson.ObjectIdHex(pid)
 	}
 
+	// 筛选-用户
+	if ownerType == constants.OwnerTypeAll {
+		user := services.GetCurrentUser(c)
+		if user.Role == constants.RoleNormal {
+			filter["$or"] = []bson.M{
+				{"user_id": services.GetCurrentUserId(c)},
+				{"is_public": true},
+			}
+		}
+	} else if ownerType == constants.OwnerTypeMe {
+		filter["user_id"] = services.GetCurrentUserId(c)
+	} else if ownerType == constants.OwnerTypePublic {
+		filter["is_public"] = true
+	}
+
 	// 排序
 	sortStr := "-_id"
 	if sortKey != "" && sortDirection != "" {
@@ -815,7 +831,7 @@ func GetSpiderStats(c *gin.Context) {
 	overview.AvgWaitDuration = overview.TotalWaitDuration / taskCount
 	overview.AvgRuntimeDuration = overview.TotalRuntimeDuration / taskCount
 
-	items, err := model.GetDailyTaskStats(bson.M{"spider_id": spider.Id})
+	items, err := model.GetDailyTaskStats(bson.M{"spider_id": spider.Id, "user_id": bson.M{"user_id": services.GetCurrentUserId(c)}})
 	if err != nil {
 		log.Errorf(err.Error())
 		HandleError(http.StatusInternalServerError, c, err)

backend/routes/stats.go

@@ -3,6 +3,7 @@ package routes
 import (
 	"crawlab/constants"
 	"crawlab/model"
+	"crawlab/services"
 	"github.com/gin-gonic/gin"
 	"github.com/globalsign/mgo/bson"
 	"net/http"
@@ -22,7 +23,7 @@ func GetHomeStats(c *gin.Context) {
 	}
 
 	// 任务总数
-	taskCount, err := model.GetTaskCount(nil)
+	taskCount, err := model.GetTaskCount(bson.M{"user_id": services.GetCurrentUserId(c)})
 	if err != nil {
 		HandleError(http.StatusInternalServerError, c, err)
 		return
@@ -36,21 +37,21 @@ func GetHomeStats(c *gin.Context) {
 	}
 
 	// 爬虫总数
-	spiderCount, err := model.GetSpiderCount()
+	spiderCount, err := model.GetSpiderCount(bson.M{"user_id": services.GetCurrentUserId(c)})
 	if err != nil {
 		HandleError(http.StatusInternalServerError, c, err)
 		return
 	}
 
 	// 定时任务数
-	scheduleCount, err := model.GetScheduleCount()
+	scheduleCount, err := model.GetScheduleCount(bson.M{"user_id": services.GetCurrentUserId(c)})
 	if err != nil {
 		HandleError(http.StatusInternalServerError, c, err)
 		return
 	}
 
 	// 每日任务数
-	items, err := model.GetDailyTaskStats(bson.M{})
+	items, err := model.GetDailyTaskStats(bson.M{"user_id": services.GetCurrentUserId(c)})
 	if err != nil {
 		HandleError(http.StatusInternalServerError, c, err)
 		return

backend/routes/task.go

@@ -47,11 +47,14 @@ func GetTaskList(c *gin.Context) {
 	if data.SpiderId != "" {
 		query["spider_id"] = bson.ObjectIdHex(data.SpiderId)
 	}
-	//新增根据任务状态获取task列表
+	// 根据任务状态获取task列表
 	if data.Status != "" {
 		query["status"] = data.Status
 	}
 
+	// 获取校验
+	query = services.GetAuthQuery(query, c)
+
 	// 获取任务列表
 	tasks, err := model.GetTaskList(query, (data.PageNum-1)*data.PageSize, data.PageSize, "-create_ts")
 	if err != nil {

backend/services/auth.go

@@ -0,0 +1,20 @@
+package services
+
+import (
+	"crawlab/constants"
+	"github.com/gin-gonic/gin"
+	"github.com/globalsign/mgo/bson"
+)
+
+func GetAuthQuery(query bson.M, c *gin.Context) bson.M {
+	user := GetCurrentUser(c)
+	if user.Role == constants.RoleAdmin {
+		// 获得所有数据
+		return query
+	} else {
+		// 只获取自己的数据
+		query["user_id"] = user.Id
+		return query
+	}
+}
+

backend/services/clean.go

@@ -0,0 +1,122 @@
+package services
+
+import (
+	"crawlab/constants"
+	"crawlab/model"
+	"github.com/apex/log"
+	"github.com/globalsign/mgo/bson"
+	"runtime/debug"
+)
+
+func InitTaskCleanUserIds() {
+	adminUser, err := GetAdminUser()
+	if err != nil {
+		log.Errorf(err.Error())
+		debug.PrintStack()
+		return
+	}
+	tasks, err := model.GetTaskList(nil, 0, constants.Infinite, "+_id")
+	if err != nil {
+		log.Errorf(err.Error())
+		debug.PrintStack()
+		return
+	}
+	for _, t := range tasks {
+		if !t.ScheduleId.Valid() {
+			t.ScheduleId = bson.ObjectIdHex(constants.ObjectIdNull)
+			if err := t.Save(); err != nil {
+				log.Errorf(err.Error())
+				debug.PrintStack()
+				continue
+			}
+		}
+
+		if !t.UserId.Valid() {
+			t.UserId = adminUser.Id
+			if err := t.Save(); err != nil {
+				log.Errorf(err.Error())
+				debug.PrintStack()
+				continue
+			}
+		}
+	}
+}
+
+func InitProjectCleanUserIds() {
+	adminUser, err := GetAdminUser()
+	if err != nil {
+		log.Errorf(err.Error())
+		debug.PrintStack()
+		return
+	}
+	projects, err := model.GetProjectList(nil, "+_id")
+	if err != nil {
+		log.Errorf(err.Error())
+		debug.PrintStack()
+		return
+	}
+	for _, p := range projects {
+		if !p.UserId.Valid() {
+			p.UserId = adminUser.Id
+			if err := p.Save(); err != nil {
+				log.Errorf(err.Error())
+				debug.PrintStack()
+				continue
+			}
+		}
+	}
+}
+
+func InitSpiderCleanUserIds() {
+	adminUser, err := GetAdminUser()
+	if err != nil {
+		log.Errorf(err.Error())
+		debug.PrintStack()
+		return
+	}
+	spiders, _ := model.GetSpiderAllList(nil)
+	for _, s := range spiders {
+		if !s.UserId.Valid() {
+			s.UserId = adminUser.Id
+			if err := s.Save(); err != nil {
+				log.Errorf(err.Error())
+				debug.PrintStack()
+				continue
+			}
+		}
+	}
+}
+
+func InitScheduleCleanUserIds() {
+	adminUser, err := GetAdminUser()
+	if err != nil {
+		log.Errorf(err.Error())
+		debug.PrintStack()
+		return
+	}
+	schedules, _ := model.GetScheduleList(nil)
+	for _, s := range schedules {
+		if !s.UserId.Valid() {
+			s.UserId = adminUser.Id
+			if err := s.Save(); err != nil {
+				log.Errorf(err.Error())
+				debug.PrintStack()
+				continue
+			}
+		}
+	}
+}
+
+func InitCleanService() error {
+	if model.IsMaster() {
+		// 清理任务UserIds
+		InitTaskCleanUserIds()
+		// 清理项目UserIds
+		InitProjectCleanUserIds()
+		// 清理爬虫UserIds
+		InitSpiderCleanUserIds()
+		// 清理定时任务UserIds
+		InitScheduleCleanUserIds()
+	}
+	return nil
+}

backend/services/spider.go

@@ -545,6 +545,9 @@ func InitSpiderService() error {
 		if err := GitCron.Start(); err != nil {
 			return err
 		}
+
+		// 清理UserId
+		InitSpiderCleanUserIds()
 	}
 
 	return nil

backend/services/user.go

@@ -120,3 +120,11 @@ func GetCurrentUser(c *gin.Context) *model.User {
 func GetCurrentUserId(c *gin.Context) bson.ObjectId {
 	return GetCurrentUser(c).Id
 }
+
+func GetAdminUser() (user *model.User, err error) {
+	u, err := model.GetUserByUsername("admin")
+	if err != nil {
+		return user, err
+	}
+	return &u, nil
+}