<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// Based on github.com/jmhodges/certificatetransparency/tools/lecsv

package main

import (
	"database/sql"
	"fmt"
	"log"
	"net/url"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"

	"github.com/go-gorp/gorp"
	"github.com/google/certificate-transparency/go"
	"github.com/google/certificate-transparency/go/client"
	"github.com/jcjones/ct-sql/censysdata"
	"github.com/jcjones/ct-sql/sqldb"
	"github.com/jcjones/ct-sql/utils"
	"github.com/jpillora/backoff"
)

var (
	config = utils.NewCTConfig()
)

type CtLogEntry struct {
	LogEntry *ct.LogEntry
	LogID    int
}

type LogDownloader struct {
	Database            *sqldb.EntriesDatabase
	EntryChan           chan CtLogEntry
	Display             *utils.ProgressDisplay
	ThreadWaitGroup     *sync.WaitGroup
	DownloaderWaitGroup *sync.WaitGroup
	Backoff             *backoff.Backoff
}

func NewLogDownloader(db *sqldb.EntriesDatabase) *LogDownloader {
	return &LogDownloader{
		Database:            db,
		EntryChan:           make(chan CtLogEntry, 25),
		Display:             utils.NewProgressDisplay(),
		ThreadWaitGroup:     new(sync.WaitGroup),
		DownloaderWaitGroup: new(sync.WaitGroup),
		Backoff: &backoff.Backoff{
			Min:    10 * time.Millisecond,
			Max:    1 * time.Second,
			Jitter: true,
		},
	}
}

func (ld *LogDownloader) StartThreads() {
	numWorkers := *config.NumThreads * runtime.NumCPU()
	for i := 0; i < numWorkers; i++ {
		go ld.insertCTWorker()
	}
}

func (ld *LogDownloader) Stop() {
	close(ld.EntryChan)
	ld.Display.Close()
}

func (ld *LogDownloader) Download(ctLogUrl string) {
	if *config.OffsetByte > 0 {
		log.Printf("[%s] Cannot set offsetByte for CT log downloads", ctLogUrl)
		return
	}

	ctLog := client.New(ctLogUrl, nil)

	log.Printf("[%s] Fetching signed tree head... ", ctLogUrl)
	sth, err := ctLog.GetSTH()
	if err != nil {
		log.Printf("[%s] Unable to fetch signed tree head: %s", ctLogUrl, err)
		return
	}

	// Set pointer in DB, now that we've verified the log works
	urlParts, err := url.Parse(ctLogUrl)
	if err != nil {
		log.Printf("[%s] Unable to parse Certificate Log: %s", ctLogUrl, err)
		return
	}
	logObj, err := ld.Database.GetLogState(fmt.Sprintf("%s%s", urlParts.Host, urlParts.Path))
	if err != nil {
		log.Printf("[%s] Unable to set Certificate Log: %s", ctLogUrl, err)
		return
	}

	var origCount uint64
	// Now we're OK to use the DB
	if *config.Offset > 0 {
		log.Printf("[%s] Starting from offset %d", ctLogUrl, *config.Offset)
		origCount = *config.Offset
	} else {
		log.Printf("[%s] Counting existing entries... ", ctLogUrl)
		origCount = logObj.MaxEntry
		if err != nil {
			log.Printf("[%s] Failed to read entries file: %s", ctLogUrl, err)
			return
		}
	}

	log.Printf("[%s] %d total entries at %s\n", ctLogUrl, sth.TreeSize, sqldb.Uint64ToTimestamp(sth.Timestamp).Format(time.ANSIC))
	if origCount == sth.TreeSize {
		log.Printf("[%s] Nothing to do\n", ctLogUrl)
		return
	}

	endPos := sth.TreeSize
	if *config.Limit > 0 && endPos > origCount+*config.Limit {
		endPos = origCount + *config.Limit
	}

	log.Printf("[%s] Going from %d to %d\n", ctLogUrl, origCount, endPos)

	finalIndex, finalTime, err := ld.DownloadCTRangeToChannel(logObj.LogID, ctLog, origCount, endPos)
	if err != nil {
		log.Printf("\n[%s] Download halting, error caught: %s\n", ctLogUrl, err)
	}

	logObj.MaxEntry = finalIndex
	if finalTime != 0 {
		logObj.LastEntryTime = sqldb.Uint64ToTimestamp(finalTime)
	}

	log.Printf("[%s] Saved state. MaxEntry=%d, LastEntryTime=%s", ctLogUrl, logObj.MaxEntry, logObj.LastEntryTime)
	ld.Database.SaveLogState(logObj)
}

// DownloadRange downloads log entries from the given starting index till one
// less than upTo. If status is not nil then status updates will be written to
// it until the function is complete, when it will be closed. The log entries
// are provided to an output channel.
func (ld *LogDownloader) DownloadCTRangeToChannel(logID int, ctLog *client.LogClient, start, upTo uint64) (uint64, uint64, error) {
	if ld.EntryChan == nil {
		return start, 0, fmt.Errorf("No output channel provided")
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, os.Interrupt)
	defer signal.Stop(sigChan)
	defer close(sigChan)

	progressTicker := time.NewTicker(10 * time.Second)
	defer progressTicker.Stop()

	var lastTime uint64

	index := start
	for index < upTo {
		max := index + 2000
		if max >= upTo {
			max = upTo - 1
		}
		rawEnts, err := ctLog.GetEntries(int64(index), int64(max))
		if err != nil {
			return index, lastTime, err
		}

		for arrayOffset := 0; arrayOffset < len(rawEnts); {
			ent := rawEnts[arrayOffset]
			// Are there waiting signals?
			select {
			case sig := <-sigChan:
				return index, lastTime, fmt.Errorf("Signal caught: %s", sig)
			case ld.EntryChan <- CtLogEntry{&ent, logID}:
				lastTime = ent.Leaf.TimestampedEntry.Timestamp
				if uint64(ent.Index) != index {
					return index, lastTime, fmt.Errorf("Index mismatch, local: %v, remote: %v", index, ent.Index)
				}

				index++
				arrayOffset++
				ld.Backoff.Reset()
			case <-progressTicker.C:
				ld.Display.UpdateProgress(fmt.Sprintf("%d", logID), start, index, upTo)
			default:
				// Channel full, retry
				time.Sleep(ld.Backoff.Duration())
			}
		}
	}

	return index, lastTime, nil
}

func (ld *LogDownloader) insertCTWorker() {
	ld.ThreadWaitGroup.Add(1)
	defer ld.ThreadWaitGroup.Done()
	for ep := range ld.EntryChan {
		err := ld.Database.InsertCTEntry(ep.LogEntry, ep.LogID)
		if err != nil {
			log.Printf("Problem inserting certificate: index: %d log: %s error: %s", ep.LogEntry.Index, *config.LogUrl, err)
		}
	}
}

func processImporter(importer censysdata.Importer, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) error {
	entryChan := make(chan censysdata.CensysEntry, 100)
	defer close(entryChan)
	wg.Add(1)
	defer wg.Done()

	display := utils.NewProgressDisplay()
	display.StartDisplay(wg)
	defer display.Close()
<commit_msg>Don't use buffered channels, as they may lose entries on shutdown<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// Based on github.com/jmhodges/certificatetransparency/tools/lecsv

package main

import (
	"database/sql"
	"fmt"
	"log"
	"net/url"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"

	"github.com/go-gorp/gorp"
	"github.com/google/certificate-transparency/go"
	"github.com/google/certificate-transparency/go/client"
	"github.com/jcjones/ct-sql/censysdata"
	"github.com/jcjones/ct-sql/sqldb"
	"github.com/jcjones/ct-sql/utils"
	"github.com/jpillora/backoff"
)

var (
	config = utils.NewCTConfig()
)

type CtLogEntry struct {
	LogEntry *ct.LogEntry
	LogID    int
}

type LogDownloader struct {
	Database            *sqldb.EntriesDatabase
	EntryChan           chan CtLogEntry
	Display             *utils.ProgressDisplay
	ThreadWaitGroup     *sync.WaitGroup
	DownloaderWaitGroup *sync.WaitGroup
	Backoff             *backoff.Backoff
}

func NewLogDownloader(db *sqldb.EntriesDatabase) *LogDownloader {
	return &LogDownloader{
		Database:            db,
		EntryChan:           make(chan CtLogEntry),
		Display:             utils.NewProgressDisplay(),
		ThreadWaitGroup:     new(sync.WaitGroup),
		DownloaderWaitGroup: new(sync.WaitGroup),
		Backoff: &backoff.Backoff{
			Min:    10 * time.Millisecond,
			Max:    1 * time.Second,
			Jitter: true,
		},
	}
}

func (ld *LogDownloader) StartThreads() {
	numWorkers := *config.NumThreads * runtime.NumCPU()
	for i := 0; i < numWorkers; i++ {
		go ld.insertCTWorker()
	}
}

func (ld *LogDownloader) Stop() {
	close(ld.EntryChan)
	ld.Display.Close()
}

func (ld *LogDownloader) Download(ctLogUrl string) {
	if *config.OffsetByte > 0 {
		log.Printf("[%s] Cannot set offsetByte for CT log downloads", ctLogUrl)
		return
	}

	ctLog := client.New(ctLogUrl, nil)

	log.Printf("[%s] Fetching signed tree head... ", ctLogUrl)
	sth, err := ctLog.GetSTH()
	if err != nil {
		log.Printf("[%s] Unable to fetch signed tree head: %s", ctLogUrl, err)
		return
	}

	// Set pointer in DB, now that we've verified the log works
	urlParts, err := url.Parse(ctLogUrl)
	if err != nil {
		log.Printf("[%s] Unable to parse Certificate Log: %s", ctLogUrl, err)
		return
	}
	logObj, err := ld.Database.GetLogState(fmt.Sprintf("%s%s", urlParts.Host, urlParts.Path))
	if err != nil {
		log.Printf("[%s] Unable to set Certificate Log: %s", ctLogUrl, err)
		return
	}

	var origCount uint64
	// Now we're OK to use the DB
	if *config.Offset > 0 {
		log.Printf("[%s] Starting from offset %d", ctLogUrl, *config.Offset)
		origCount = *config.Offset
	} else {
		log.Printf("[%s] Counting existing entries... ", ctLogUrl)
		origCount = logObj.MaxEntry
		if err != nil {
			log.Printf("[%s] Failed to read entries file: %s", ctLogUrl, err)
			return
		}
	}

	log.Printf("[%s] %d total entries at %s\n", ctLogUrl, sth.TreeSize, sqldb.Uint64ToTimestamp(sth.Timestamp).Format(time.ANSIC))
	if origCount == sth.TreeSize {
		log.Printf("[%s] Nothing to do\n", ctLogUrl)
		return
	}

	endPos := sth.TreeSize
	if *config.Limit > 0 && endPos > origCount+*config.Limit {
		endPos = origCount + *config.Limit
	}

	log.Printf("[%s] Going from %d to %d\n", ctLogUrl, origCount, endPos)

	finalIndex, finalTime, err := ld.DownloadCTRangeToChannel(logObj.LogID, ctLog, origCount, endPos)
	if err != nil {
		log.Printf("\n[%s] Download halting, error caught: %s\n", ctLogUrl, err)
	}

	logObj.MaxEntry = finalIndex
	if finalTime != 0 {
		logObj.LastEntryTime = sqldb.Uint64ToTimestamp(finalTime)
	}

	log.Printf("[%s] Saved state. MaxEntry=%d, LastEntryTime=%s", ctLogUrl, logObj.MaxEntry, logObj.LastEntryTime)
	ld.Database.SaveLogState(logObj)
}

// DownloadRange downloads log entries from the given starting index till one
// less than upTo. If status is not nil then status updates will be written to
// it until the function is complete, when it will be closed. The log entries
// are provided to an output channel.
func (ld *LogDownloader) DownloadCTRangeToChannel(logID int, ctLog *client.LogClient, start, upTo uint64) (uint64, uint64, error) {
	if ld.EntryChan == nil {
		return start, 0, fmt.Errorf("No output channel provided")
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, os.Interrupt)
	defer signal.Stop(sigChan)
	defer close(sigChan)

	progressTicker := time.NewTicker(10 * time.Second)
	defer progressTicker.Stop()

	var lastTime uint64

	index := start
	for index < upTo {
		max := index + 2000
		if max >= upTo {
			max = upTo - 1
		}
		rawEnts, err := ctLog.GetEntries(int64(index), int64(max))
		if err != nil {
			return index, lastTime, err
		}

		for arrayOffset := 0; arrayOffset < len(rawEnts); {
			ent := rawEnts[arrayOffset]
			// Are there waiting signals?
			select {
			case sig := <-sigChan:
				return index, lastTime, fmt.Errorf("Signal caught: %s", sig)
			case ld.EntryChan <- CtLogEntry{&ent, logID}:
				lastTime = ent.Leaf.TimestampedEntry.Timestamp
				if uint64(ent.Index) != index {
					return index, lastTime, fmt.Errorf("Index mismatch, local: %v, remote: %v", index, ent.Index)
				}

				index++
				arrayOffset++
				ld.Backoff.Reset()
			case <-progressTicker.C:
				ld.Display.UpdateProgress(fmt.Sprintf("%d", logID), start, index, upTo)
			default:
				// Channel full, retry
				time.Sleep(ld.Backoff.Duration())
			}
		}
	}

	return index, lastTime, nil
}

func (ld *LogDownloader) insertCTWorker() {
	ld.ThreadWaitGroup.Add(1)
	defer ld.ThreadWaitGroup.Done()
	for ep := range ld.EntryChan {
		err := ld.Database.InsertCTEntry(ep.LogEntry, ep.LogID)
		if err != nil {
			log.Printf("Problem inserting certificate: index: %d log: %s error: %s", ep.LogEntry.Index, *config.LogUrl, err)
		}
	}
}

func processImporter(importer censysdata.Importer, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) error {
	entryChan := make(chan censysdata.CensysEntry)
	defer close(entryChan)
	wg.Add(1)
	defer wg.Done()

	display := utils.NewProgressDisplay()
	display.StartDisplay(wg)
	defer display.Close()

{'fraction_non_alphanumeric': 0.13347337803991274, 'fraction_numerical': 0.007513937141472085, 'mean_word_length': 3.7262313860252005, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 36, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3347849', 'n_tokens_mistral': 4825, 'n_tokens_neox': 4414, 'n_words': 1136}