feat(Enterprise version supporting training camps)

cmd/root.go

@@ -4,7 +4,6 @@ import (
  "context"
  "errors"
  "fmt"
- "io/ioutil"
  "math"
  "math/rand"
  "net/http"
@@ -30,22 +29,23 @@ import (
 )
 
 var (
- phone string
- gcid string
- gcess string
- concurrency int
- downloadFolder string
- sp *spinner.Spinner
- selectedProduct geektime.Product
- quality string
- downloadComments bool
- selectedProductType productTypeSelectOption
- columnOutputType int
- waitSeconds int
- productTypeOptions = make([]productTypeSelectOption, 6)
- geektimeClient *geektime.Client
- accountClient *geektime.Client
- universityClient *geektime.Client
+ phone string
+ gcid string
+ gcess string
+ concurrency int
+ downloadFolder string
+ sp *spinner.Spinner
+ selectedProduct geektime.Product
+ quality string
+ downloadComments bool
+ selectedProductType productTypeSelectOption
+ columnOutputType int
+ waitSeconds int
+ productTypeOptions = make([]productTypeSelectOption, 7)
+ geektimeClient *geektime.Client
+ geekEnterpriseClient *geektime.Client
+ accountClient *geektime.Client
+ universityClient *geektime.Client
 )
 
 type productTypeSelectOption struct {
@@ -90,6 +90,7 @@ func setProductTypeOptions() {
  productTypeOptions[3] = productTypeSelectOption{3, "大厂案例", 4, []string{"q"}, false}
  productTypeOptions[4] = productTypeSelectOption{4, "训练营", 5, []string{""}, true} //custom source type, not use
  productTypeOptions[5] = productTypeSelectOption{5, "其他", 1, []string{"x", "c6"}, true}
+ productTypeOptions[6] = productTypeSelectOption{6, "企业版训练营", 6, []string{"c44"}, true}
 }
 
 var rootCmd = &cobra.Command{
@@ -145,6 +146,7 @@ var rootCmd = &cobra.Command{
  }
  geektimeClient = geektime.NewClient(readCookies)
  universityClient = geektime.NewUniversityClient(readCookies)
+ geekEnterpriseClient = geektime.NewEnterpriseClient(readCookies)
  selectProductType(cmd.Context())
  },
 }
@@ -232,6 +234,8 @@ func loadProduct(ctx context.Context, productID int) {
  p, err = universityClient.MyClassProduct(productID)
  // university don't need check product type
  // if input invalid id, access mark is 0
+ } else if isEnterpriseUniversity() {
+ p, err = geekEnterpriseClient.EnterpriseArticlesInfo(productID)
  } else {
  p, err = geektimeClient.ColumnInfo(productID)
  if err == nil {
@@ -419,15 +423,24 @@ func handleDownloadAll(ctx context.Context) {
  }
  } else {
  for _, a := range selectedProduct.Articles {
+ sectionDir := projectDir
  fileName := filenamify.Filenamify(a.Title) + video.TSExtension
  if _, ok := downloaded[fileName]; ok {
  continue
  }
+ // add sub dir
+ if a.SectionTitle != "" {
+ sectionDir, err = mkDownloadProjectSectionDir(projectDir, a.SectionTitle)
+ checkError(err)
+ }
  if isUniversity() {
- err := video.DownloadUniversityVideo(ctx, universityClient, a.AID, selectedProduct, projectDir, quality, concurrency)
+ err := video.DownloadUniversityVideo(ctx, universityClient, a.AID, selectedProduct, sectionDir, quality, concurrency)
+ checkError(err)
+ } else if isEnterpriseUniversity() {
+ err := video.DownloadEnterpriseArticleVideo(ctx, geekEnterpriseClient, a.AID, selectedProductType.SourceType, sectionDir, quality, concurrency)
  checkError(err)
  } else {
- err := video.DownloadArticleVideo(ctx, geektimeClient, a.AID, selectedProductType.SourceType, projectDir, quality, concurrency)
+ err := video.DownloadArticleVideo(ctx, geektimeClient, a.AID, selectedProductType.SourceType, sectionDir, quality, concurrency)
  checkError(err)
  }
  }
@@ -441,7 +454,7 @@ func increasePDFCount(total int, i *int) {
 }
 
 func loadArticles() {
- if !isUniversity() && len(selectedProduct.Articles) <= 0 {
+ if !isUniversity() && !isEnterpriseUniversity() && len(selectedProduct.Articles) <= 0 {
  sp.Prefix = "[ 正在加载文章列表... ]"
  sp.Start()
  articles, err := geektimeClient.ColumnArticles(strconv.Itoa(selectedProduct.ID))
@@ -513,6 +526,9 @@ func downloadArticle(ctx context.Context, article geektime.Article, projectDir s
  if isUniversity() {
  err := video.DownloadUniversityVideo(ctx, universityClient, article.AID, selectedProduct, projectDir, quality, concurrency)
  checkError(err)
+ } else if isEnterpriseUniversity() {
+ err := video.DownloadEnterpriseArticleVideo(ctx, geekEnterpriseClient, article.AID, selectedProductType.SourceType, projectDir, quality, concurrency)
+ checkError(err)
  } else {
  err := video.DownloadArticleVideo(ctx, geektimeClient, article.AID, selectedProductType.SourceType, projectDir, quality, concurrency)
  checkError(err)
@@ -528,6 +544,10 @@ func isUniversity() bool {
  return selectedProductType.Index == 4
 }
 
+func isEnterpriseUniversity() bool {
+ return selectedProductType.Index == 6
+}
+
 // Sets the bit at pos in the integer n.
 func setBit(n int, pos uint) int {
  n |= (1 << pos)
@@ -555,17 +575,24 @@ func readCookiesFromInput() []*http.Cookie {
 }
 
 func findDownloadedArticleFileNames(projectDir string) (map[string]struct{}, error) {
- files, err := ioutil.ReadDir(projectDir)
- res := make(map[string]struct{}, len(files))
- if err != nil {
- return res, err
- }
- if len(files) == 0 {
- return res, nil
- }
- for _, f := range files {
- res[f.Name()] = struct{}{}
- }
+ res := make(map[string]struct{})
+ limit := 2
+ err := filepath.Walk(projectDir, func(path string, info os.FileInfo, err error) error {
+ if err != nil {
+ fmt.Printf("访问路径时出错：%v\n", err)
+ return err
+ }
+ // 计算当前路径的深度
+ depth := len(filepath.SplitList(path)) - len(filepath.SplitList(projectDir))
+ if depth >= limit {
+ return filepath.SkipDir // 如果达到限制深度，则跳过该文件夹及其子文件夹
+ }
+ if !info.IsDir() {
+ res[info.Name()] = struct{}{}
+ }
+ return nil
+ })
+ checkError(err)
  return res, nil
 }
 
@@ -582,6 +609,15 @@ func mkDownloadProjectDir(downloadFolder, phone, gcid, projectName string) (stri
  return path, nil
 }
 
+func mkDownloadProjectSectionDir(downloadFolder, sectionName string) (string, error) {
+ path := filepath.Join(downloadFolder, filenamify.Filenamify(sectionName))
+ err := os.MkdirAll(path, os.ModePerm)
+ if err != nil {
+ return "", err
+ }
+ return path, nil
+}
+
 func checkProductType(productType string) bool {
  for _, pt := range selectedProductType.AcceptProductTypes {
  if pt == productType {

internal/geektime/geektime.go

@@ -24,6 +24,7 @@ const (
  UserAgent = "User-Agent"
  // GeekBangUniversityBaseURL ...
  GeekBangUniversityBaseURL = "https://u.geekbang.org"
+ GeekBangEnterpriseBaseURL = "https://b.geekbang.org"
  // GeekBangAccountBaseURL ...
  GeekBangAccountBaseURL = "https://account.geekbang.org"
  // LoginPath ...
@@ -49,6 +50,15 @@ const (
  // UniversityV1MyClassInfoPath get university class info and all articles info in it
  UniversityV1MyClassInfoPath = "/serv/v1/myclass/info"
 
+ // V1EnterpriseProductInfoPath used in enterprise course product info
+ V1EnterpriseProductInfoPath = "/app/v1/course/info"
+ // V1EnterpriseArticlesInfoPath used in enterprise course articles info
+ V1EnterpriseArticlesInfoPath = "/app/v1/course/articles"
+ // V1EnterpriseArticleDetailInfoPath used in enterprise course article detail info
+ V1EnterpriseArticleDetailInfoPath = "/app/v1/article/detail"
+ // V1EnterpriseVideoPlayAuthPath used in enterprise course video play auth
+ V1EnterpriseVideoPlayAuthPath = "/app/v1/source_auth/video_play_auth"
+
  // GeekBangCookieDomain ...
  GeekBangCookieDomain = ".geekbang.org"
 
@@ -81,8 +91,9 @@ type Product struct {
 
 // Article ...
 type Article struct {
- AID int
- Title string
+ AID int
+ SectionTitle string
+ Title string
 }
 
 // ErrGeekTimeAPIBadCode ...
@@ -141,6 +152,19 @@ func NewUniversityClient(cs []*http.Cookie) *Client {
  return c
 }
 
+// NewEnterpriseClient
+func NewEnterpriseClient(cs []*http.Cookie) *Client {
+ httpClient := resty.New().
+ SetCookies(cs).
+ SetRetryCount(1).
+ SetTimeout(10*time.Second).
+ SetHeader("User-Agent", DefaultUserAgent).
+ SetLogger(logger.DiscardLogger{})
+
+ c := &Client{HTTPClient: httpClient, BaseURL: GeekBangEnterpriseBaseURL, Cookies: cs}
+ return c
+}
+
 // Login call geektime login api and return auth cookies
 func (c *Client) Login(phone, password string) ([]*http.Cookie, error) {
  var res struct {
@@ -328,6 +352,109 @@ func (c *Client) ProductInfo(productID int) (response.V3ProductInfoResponse, err
  return res, nil
 }
 
+func (c *Client) enterpriseProductInfo(productID int) (response.V1EnterpriseProductInfoResponse, error) {
+ var res response.V1EnterpriseProductInfoResponse
+ r := c.newRequest(resty.MethodPost,
+ V1EnterpriseProductInfoPath,
+ nil,
+ map[string]interface{}{
+ "id": productID,
+ },
+ &res,
+ )
+ if _, err := do(r); err != nil {
+ return response.V1EnterpriseProductInfoResponse{}, err
+ }
+ return res, nil
+}
+
+func (c *Client) EnterpriseArticlesInfo(id int) (Product, error) {
+ var p Product
+ productInfo, err := c.enterpriseProductInfo(id)
+ if err != nil {
+ return p, err
+ }
+
+ var res response.V1EnterpriseArticlesResponse
+ r := c.newRequest(resty.MethodPost,
+ V1EnterpriseArticlesInfoPath,
+ nil,
+ map[string]interface{}{
+ "id": id,
+ },
+ &res,
+ )
+
+ resp, err := do(r)
+ if err != nil {
+ return p, err
+ }
+
+ if res.Code != 0 {
+ if !res.Data.IsShow && !productInfo.Data.Extra.IsMyCourse {
+ p.Access = false
+ return p, nil
+ }
+ return p, ErrGeekTimeAPIBadCode{V1EnterpriseArticlesInfoPath, resp.String()}
+ }
+
+ p = Product{
+ Access: true,
+ ID: id,
+ Title: productInfo.Data.Title,
+ Type: "",
+ IsVideo: true,
+ }
+ var articles []Article
+
+ for _, sections := range res.Data.List {
+ for _, a := range sections.ArticleList {
+ articleID, _ := strconv.Atoi(a.Article.ID)
+ articles = append(articles, Article{
+ AID: articleID,
+ SectionTitle: sections.Title,
+ Title: a.Article.Title,
+ })
+ }
+ }
+ p.Articles = articles
+
+ return p, nil
+}
+
+func (c *Client) V1EnterpriseArticleDetailInfo(articleID string) (response.V1EnterpriseArticlesDetailResponse, error) {
+ var res response.V1EnterpriseArticlesDetailResponse
+ r := c.newRequest(resty.MethodPost,
+ V1EnterpriseArticleDetailInfoPath,
+ nil,
+ map[string]interface{}{
+ "article_id": articleID,
+ },
+ &res,
+ )
+ if _, err := do(r); err != nil {
+ return response.V1EnterpriseArticlesDetailResponse{}, err
+ }
+ return res, nil
+}
+
+func (c *Client) EnterpriseVideoPlayAuth(articleID, videoID string) (string, error) {
+ var res response.V3VideoPlayAuthResponse
+ r := c.newRequest(resty.MethodPost,
+ V1EnterpriseVideoPlayAuthPath,
+ nil,
+ map[string]interface{}{
+ "aid": articleID,
+ "video_id": videoID,
+ },
+ &res,
+ )
+ if _, err := do(r); err != nil {
+ return "", err
+ }
+ return res.Data.PlayAuth, nil
+}
+
 // V3ArticleInfo used to get daily lesson or qconplus article info
 func (c *Client) V3ArticleInfo(articleID int) (response.V3ArticleInfoResponse, error) {
  var res response.V3ArticleInfoResponse