|
1 | 1 | package main |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "archive/zip" |
4 | 5 | "flag" |
| 6 | + "io" |
5 | 7 | "log" |
| 8 | + "net/http" |
| 9 | + "os" |
| 10 | + "path/filepath" |
6 | 11 | "strings" |
7 | 12 |
|
8 | 13 | brain "github.com/sundy-li/wechat_brain" |
| 14 | + |
| 15 | + "github.com/PuerkitoBio/goquery" |
9 | 16 | ) |
10 | 17 |
|
11 | 18 | var ( |
12 | | - action string |
13 | | - fs string |
| 19 | + source string |
| 20 | + fs string |
| 21 | + issueUrl = "https://github.com/sundy-li/wechat_brain/issues/17" |
| 22 | + tmpDir = "/data/tmp/" |
14 | 23 | ) |
15 | 24 |
|
16 | 25 | func init() { |
17 | | - flag.StringVar(&action, "a", "show", "action value -> show | merge") |
| 26 | + flag.StringVar(&source, "s", "show", "source value -> show | file | issue") |
18 | 27 | flag.StringVar(&fs, "fs", "", "merge data files") |
19 | 28 | flag.Parse() |
20 | 29 | } |
21 | 30 |
|
22 | 31 | func main() { |
23 | | - if action == "merge" { |
| 32 | + if source == "file" { |
24 | 33 | files := strings.Split(fs, " ") |
25 | 34 | if len(files) < 1 { |
26 | 35 | log.Println("empty files") |
27 | 36 | return |
28 | 37 | } |
29 | 38 | brain.MergeQuestions(files...) |
| 39 | + } else if source == "issue" { |
| 40 | + doc, _ := goquery.NewDocument(issueUrl) |
| 41 | + doc.Find("div.comment").Each(func(index int, comment *goquery.Selection) { |
| 42 | + comment.Find("td.d-block p a").Each(func(i int, s *goquery.Selection) { |
| 43 | + if strings.Contains(s.Text(), "questions.zip") { |
| 44 | + href, _ := s.Attr("href") |
| 45 | + if href != "" { |
| 46 | + handleZipUrl(href) |
| 47 | + } |
| 48 | + } |
| 49 | + }) |
| 50 | + }) |
30 | 51 | } |
31 | 52 | total := brain.CountQuestions() |
32 | 53 | log.Println("total questions =>", total) |
33 | 54 | } |
| 55 | + |
| 56 | +func handleZipUrl(url string) error { |
| 57 | + resp, err := http.Get(url) |
| 58 | + if err != nil { |
| 59 | + return err |
| 60 | + } |
| 61 | + defer resp.Body.Close() |
| 62 | + out, err := os.Create(tmpDir + "questions.zip") |
| 63 | + if err != nil { |
| 64 | + return err |
| 65 | + } |
| 66 | + defer out.Close() |
| 67 | + io.Copy(out, resp.Body) |
| 68 | + _, err = Unzip(tmpDir+"questions.zip", tmpDir) |
| 69 | + if err != nil { |
| 70 | + return err |
| 71 | + } |
| 72 | + |
| 73 | + //merge data |
| 74 | + brain.MergeQuestions(tmpDir + "questions.data") |
| 75 | + log.Println("merged", url) |
| 76 | + return nil |
| 77 | +} |
| 78 | + |
| 79 | +// Unzip will un-compress a zip archive, |
| 80 | +// moving all files and folders to an output directory |
| 81 | +func Unzip(src, dest string) ([]string, error) { |
| 82 | + |
| 83 | + var filenames []string |
| 84 | + |
| 85 | + r, err := zip.OpenReader(src) |
| 86 | + if err != nil { |
| 87 | + return filenames, err |
| 88 | + } |
| 89 | + defer r.Close() |
| 90 | + |
| 91 | + for _, f := range r.File { |
| 92 | + |
| 93 | + rc, err := f.Open() |
| 94 | + if err != nil { |
| 95 | + return filenames, err |
| 96 | + } |
| 97 | + defer rc.Close() |
| 98 | + |
| 99 | + // Store filename/path for returning and using later on |
| 100 | + fpath := filepath.Join(dest, f.Name) |
| 101 | + filenames = append(filenames, fpath) |
| 102 | + |
| 103 | + if f.FileInfo().IsDir() { |
| 104 | + // Make Folder |
| 105 | + os.MkdirAll(fpath, os.ModePerm) |
| 106 | + |
| 107 | + } else { |
| 108 | + |
| 109 | + // Make File |
| 110 | + var fdir string |
| 111 | + if lastIndex := strings.LastIndex(fpath, string(os.PathSeparator)); lastIndex > -1 { |
| 112 | + fdir = fpath[:lastIndex] |
| 113 | + } |
| 114 | + |
| 115 | + err = os.MkdirAll(fdir, os.ModePerm) |
| 116 | + if err != nil { |
| 117 | + log.Fatal(err) |
| 118 | + return filenames, err |
| 119 | + } |
| 120 | + f, err := os.OpenFile( |
| 121 | + fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode()) |
| 122 | + if err != nil { |
| 123 | + return filenames, err |
| 124 | + } |
| 125 | + defer f.Close() |
| 126 | + |
| 127 | + _, err = io.Copy(f, rc) |
| 128 | + if err != nil { |
| 129 | + return filenames, err |
| 130 | + } |
| 131 | + |
| 132 | + } |
| 133 | + } |
| 134 | + return filenames, nil |
| 135 | +} |
0 commit comments