/
cloud.go
110 lines (97 loc) · 2.95 KB
/
cloud.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
package collect
import (
"context"
"errors"
"strings"
"time"
"github.com/jyothri/hdd/constants"
"github.com/jyothri/hdd/db"
"golang.org/x/oauth2"
"golang.org/x/oauth2/google"
"google.golang.org/api/drive/v3"
"google.golang.org/api/googleapi"
"google.golang.org/api/option"
)
// List of fields to be retreived on file resource from the drive API.
var fields []string = []string{"size", "id", "name", "mimeType", "parents", "modifiedTime", "md5Checksum"}
var paginationFields []string = []string{"nextPageToken", "incompleteSearch"}
const pageSize = 1000
var cloudConfig *oauth2.Config
func init() {
cloudConfig = &oauth2.Config{
ClientID: constants.OauthClientId,
ClientSecret: constants.OauthClientSecret,
Endpoint: google.Endpoint,
Scopes: []string{drive.DriveReadonlyScope},
}
}
func getDriveService(refreshToken string) *drive.Service {
tokenSrc := oauth2.Token{
RefreshToken: refreshToken,
}
ctx := context.Background()
driveService, err := drive.NewService(ctx, option.WithTokenSource(cloudConfig.TokenSource(ctx, &tokenSrc)))
checkError(err)
return driveService
}
func CloudDrive(driveScan GDriveScan) int {
scanData := make(chan db.FileData, 10)
scanId := db.LogStartScan("google_drive")
driveService := getDriveService(driveScan.RefreshToken)
go db.SaveScanMetadata("", driveScan.QueryString, scanId)
go startCloudDrive(driveService, scanId, driveScan.QueryString, scanData)
go db.SaveStatToDb(scanId, scanData)
return scanId
}
func startCloudDrive(driveService *drive.Service, scanId int, queryString string, scanData chan<- db.FileData) {
lock.Lock()
defer lock.Unlock()
filesListCall := driveService.Files.List().PageSize(pageSize).Q(queryString).Fields(googleapi.Field(strings.Join(append(addPrefix(fields, "files/"), paginationFields...), ",")))
hasNextPage := true
for hasNextPage {
fileList, err := filesListCall.Do()
checkError(err)
if fileList.IncompleteSearch {
checkError(errors.New("incomplete search from drive API"))
}
parseFileList(fileList, scanData)
if fileList.NextPageToken == "" {
hasNextPage = false
}
filesListCall = filesListCall.PageToken(fileList.NextPageToken)
}
close(scanData)
}
func parseFileList(fileList *drive.FileList, scanData chan<- db.FileData) {
for _, file := range fileList.Files {
fd := db.FileData{
FileName: file.Name,
FilePath: file.Id,
IsDir: file.MimeType == "application/vnd.google-apps.folder",
ModTime: parseTime(file.ModifiedTime),
FileCount: 1,
}
if !fd.IsDir {
fd.Size = uint(file.Size)
fd.FileCount = 1
fd.Md5Hash = file.Md5Checksum
scanData <- fd
}
}
}
func addPrefix(in []string, prefix string) []string {
out := make([]string, len(in))
for idx, str := range in {
out[idx] = prefix + str
}
return out
}
func parseTime(inputTime string) time.Time {
parsedTime, err := time.Parse(time.RFC3339, inputTime)
checkError(err)
return parsedTime
}
type GDriveScan struct {
QueryString string
RefreshToken string
}