main.go

// This file is part of GoCrawl.
//
// GoCrawl is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation, either version 3 of the License, or (at your
// option) any later version.
//
// GoCrawl is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
// for more details.
//
// You should have received a copy of the GNU General Public License along
// with GoCrawl. If not, see <https://www.gnu.org/licenses/>.

package main

import (
	"context"
	"database/sql"
	"fmt"
	"github.com/namsral/flag"
	"io"
	"log"
	"os"
	"time"

	_ "github.com/mattn/go-sqlite3"

	"gocrawl/internal/crawler"
	"gocrawl/internal/graph"
)

type Config struct {
	UserAgent   string
	ContentType string
	Url         string
	MaxWorkers  int
}

const defaultUserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
	"AppleWebKit/537.36 (KHTML, like Gecko) " +
	"Chrome/104.0.5112.79 Safari/537.36"

func (c *Config) Init(args []string) error {
	flags := flag.NewFlagSet(args[0], flag.ExitOnError)
	flags.String(flag.DefaultConfigFlagname, "", "Path to config file")

	var (
		userAgent   = flags.String("user_agent", defaultUserAgent, "User-Agent HTTP header value")
		contentType = flags.String("content_type", "text/html", "Content-Type HTTP header value")
		url         = flags.String("url", "", "URL to crawl if not daemon")
		maxWorkers  = flags.Int("max_workers", 8, "maximum number of goroutines")
		tick        = flags.Duration("tick", defaultTick, "Polling of worker goroutines before crawling")
	)

	if err := flags.Parse(args[1:]); err != nil {
		return err
	}

	c.UserAgent = *userAgent
	c.ContentType = *contentType
	c.Url = *url
	c.MaxWorkers = *maxWorkers

	return nil
}

func main() {
	c := &Config{}
	c.Init(os.Args)

	ctx := context.WithValue(context.Background(), "config", c)
	ctx, cancel := context.WithCancel(ctx)
	defer func() {
		cancel()
	}()

	if err := run(ctx, os.Stdout); err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		os.Exit(1)
	}
}

func initDatabase(db *sql.DB) error {
	db, err := sql.Open("sqlite3", "./gocrawl.db")
	if err != nil {
		return err
	}

	if err = graph.InitDbFile(db); err != nil {
		return err
	}

	return nil
}

func run(ctx context.Context, out io.Writer) error {
	log.SetOutput(out)
	var db *sql.DB = nil
	if err := initDatabase(&db); err != nil {
		return err
	}
	defer db.Close()

	var pageQueue = make(chan *Page, 100)

	log.Printf("Spawning %d workers", ctx.Value("config").MaxWorkers)
	for i := 0; i < ctx.Value("config").MaxWorkers; i++ {
		go crawler.Bot(ctx, db, pageQueue)
	}
}