示例#1
0
文件: main.go 项目: mpdroog/sla
func main() {
	var e error
	var verbose, skipyenc bool
	var configPath, date string
	flag.BoolVar(&verbose, "v", false, "Verbosity")
	flag.BoolVar(&skipyenc, "y", false, "Skip yEnc decode")
	flag.StringVar(&configPath, "c", "./config.json", "/Path/to/config.json")
	flag.StringVar(&date, "d", "", "YYYY-mm-dd to download from nzbdir")
	flag.Parse()

	C, e = loadConfig(configPath)
	if e != nil {
		fail(e)
	}
	if !strings.HasSuffix(C.NzbDir, "/") {
		C.NzbDir += "/"
	}
	if date == "" {
		// default to today
		date = time.Now().Format("2006-01-02")
	}
	if verbose {
		fmt.Printf("Config=%+v Date=%+v\n", C, date)
	}

	// Force valid date pattern
	if _, e := time.Parse("2006-01-02", date); e != nil {
		fail(e)
	}

	fd, e := os.Open(C.NzbDir + date + ".nzb")
	if e != nil {
		fail(e)
	}
	defer fd.Close()
	arts, e := nzb.Read(fd)
	if e != nil {
		fail(e)
	}

	if verbose {
		fmt.Println("Connecting to nntp..")
	}
	conn := nntp.New(C.Address, "1", verbose)
	conn.Verbose = verbose
	perfBegin := time.Now()
	var perfInit, perfAuth time.Time
	{
		defer conn.Close()
		if e := conn.Init(); e != nil {
			fail(e)
		}
		perfInit = time.Now()
		if e := conn.Auth(C.User, C.Pass); e != nil {
			fail(e)
		}
		perfAuth = time.Now()
	}

	perfArts := []float64{}
	KBsecs := []float64{}
	lastPerf := time.Now()
	buf := new(bytes.Buffer)
	for _, segment := range arts.File.Segments.Segment {
		buf.Reset()
		conn.Article(segment.Msgid)

		counter := stream.NewCountReader(conn.GetReader())
		rawread := bufio.NewReader(counter)

		_, e = textproto.NewReader(rawread).ReadMIMEHeader()
		if e != nil {
			fail(e)
		}

		if _, e := io.Copy(buf, rawread); e != nil {
			fail(e)
		}
		n := counter.ReadReset()
		if int64(n) <= segment.Bytes {
			panic(fmt.Errorf("ByteCount mismatch, expect>%d recv=%d", segment.Bytes, n))
		}

		if !skipyenc {
			if _, e = yenc.Decode(buf); e != nil {
				fmt.Printf("%+v\n", string(buf.Bytes()))
				fail(e)
			}
		}

		now := time.Now()
		diff := now.Sub(lastPerf)
		kbSec := float64(n/1024) / diff.Seconds()

		if verbose {
			fmt.Println(fmt.Sprintf(
				"Download %s (%d bytes in %s with %f KB/s)",
				segment.Msgid, n, diff.String(), kbSec,
			))
		}

		KBsecs = append(KBsecs, kbSec)
		perfArts = append(perfArts, duration.MilliSeconds(diff))
		lastPerf = now
	}

	enc := json.NewEncoder(os.Stdout)
	if e := enc.Encode(Perf{
		Conn:  duration.MilliSeconds(perfInit.Sub(perfBegin)),
		Auth:  duration.MilliSeconds(perfAuth.Sub(perfInit)),
		Arts:  perfArts,
		KBsec: KBsecs,
		Error: []string{},
	}); e != nil {
		fail(e)
	}
}
示例#2
0
文件: main.go 项目: mpdroog/sla
func main() {
	var verbose bool
	var configPath string

	flag.BoolVar(&verbose, "v", false, "Verbosity")
	flag.StringVar(&configPath, "c", "./config.json", "/Path/to/config.json")
	flag.Parse()

	c, e := loadConfig(configPath)
	if e != nil {
		fail(e)
	}
	if !strings.HasSuffix(c.NzbDir, "/") {
		c.NzbDir += "/"
	}

	// Permission check nzbdir
	{
		stat, e := os.Stat(c.NzbDir)
		if e != nil {
			fail(e)
		}
		if !stat.IsDir() {
			fmt.Println("Not a dir: " + c.NzbDir)
			os.Exit(1)
			return
		}
		if e := ioutil.WriteFile(
			c.NzbDir+"check.txt",
			[]byte("Write permission check."),
			0400,
		); e != nil {
			fail(e)
		}
		if e := os.Remove(c.NzbDir + "check.txt"); e != nil {
			fail(e)
		}
	}
	// Permission check uploaddir
	{
		stat, e := os.Stat(c.UploadDir)
		if e != nil {
			fail(e)
		}
		if !stat.IsDir() {
			fmt.Println("Not a dir: " + c.UploadDir)
			os.Exit(1)
			return
		}
	}

	if verbose {
		fmt.Println("Building ZIP from dir=" + c.UploadDir)
	}

	enc := yenc.NewWriter(
		new(bytes.Buffer),
		fmt.Sprintf("sla-%s.zip", time.Now().Format("2006-01-02")),
		yenc.PART_SIZE,
	)
	var partCount = 0
	{
		buf := new(bytes.Buffer)
		w := zip.NewWriter(buf)

		e := filepath.Walk(c.UploadDir, func(path string, info os.FileInfo, err error) error {
			if path == c.UploadDir {
				// Ignore base
				return nil
			}
			if strings.HasSuffix(info.Name(), ".sh") {
				// Ignore scripts
				if verbose {
					fmt.Println("Skip " + path)
				}
				return nil
			}
			if verbose {
				fmt.Println("Add " + path + " to ZIP.")
			}
			if e := zipAdd(w, info.Name(), path); e != nil {
				return e
			}
			return nil
		})
		if e != nil {
			fail(e)
		}

		f, e := w.Create("unique.txt")
		if e != nil {
			fail(e)
		}
		f.Write([]byte(RandStringRunes(16)))

		if e := w.Close(); e != nil {
			fail(e)
		}

		if _, err := enc.Write(buf.Bytes()); err != nil {
			fail(err)
		}
		partCount = enc.Parts()
	}
	if partCount < 50 {
		fmt.Printf("Need at least 50 parts, I got: %d (increase rand file?)\n", partCount)
		os.Exit(1)
	}

	subject := "Completion test " + time.Now().Format("2006-01-02")
	if verbose {
		fmt.Println(fmt.Sprintf("Upload file=%s parts(%d)..", subject, partCount))
	}

	if verbose {
		fmt.Println("Connecting to nntp..")
	}
	conn := nntp.New(c.Address, "1", verbose)
	conn.Verbose = verbose
	perfBegin := time.Now()
	var perfInit, perfAuth time.Time
	{
		defer conn.Close()
		if e := conn.Init(); e != nil {
			fail(e)
		}
		perfInit = time.Now()
		if e := conn.Auth(c.User, c.Pass); e != nil {
			fail(e)
		}
		perfAuth = time.Now()
	}

	var msgids []nzb.Msg
	artPerf := []ArtPerf{}
	lastPerf := time.Now()

	for enc.HasNext() {
		if e := conn.Post(); e != nil {
			fail(e)
		}
		msgid := RandStringRunes(16) + c.MsgDomain

		w := stream.NewCountWriter(conn.GetWriter())
		if _, e := w.WriteString(headers(subject, msgid)); e != nil {
			fail(e)
		}
		w.ResetWritten()

		if _, e := enc.EncodePart(w); e != nil {
			fail(e)
		}
		n := w.Written()
		msgids = append(msgids, nzb.Msg{
			Msgid: msgid,
			Size:  n,
		})

		if e := conn.PostClose(); e != nil {
			fail(e)
		}

		// Stats
		now := time.Now()
		d := now.Sub(lastPerf)

		if verbose {
			fmt.Println(fmt.Sprintf(
				"Posted %s in %s",
				msgid, d.String(),
			))
		}
		kbSec := float64(n/1024) / d.Seconds()
		artPerf = append(artPerf, ArtPerf{
			MsgId:    msgid,
			Time:     duration.MilliSeconds(d),
			Size:     n,
			Speed:    kbSec,     // kb/sec
			BitSpeed: kbSec * 8, // kbit/sec
		})
		lastPerf = now
	}
	if err := enc.Close(); err != nil {
		fail(err)
	}

	xml := nzb.Build(subject, msgids, time.Now().Format(time.RFC822))
	if e := ioutil.WriteFile(
		c.NzbDir+time.Now().Format("2006-01-02")+".nzb",
		[]byte(xml), 400,
	); e != nil {
		fail(e)
	}

	jenc := json.NewEncoder(os.Stdout)
	if e := jenc.Encode(Perf{
		Conn:  duration.MilliSeconds(perfInit.Sub(perfBegin)),
		Auth:  duration.MilliSeconds(perfAuth.Sub(perfInit)),
		Arts:  artPerf,
		Error: []string{},
	}); e != nil {
		fail(e)
	}
}