<commit_before>package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/asticode/go-astilog"
	"github.com/asticode/go-astitools/flag"
	"github.com/asticode/go-astits"
	"github.com/pkg/errors"
	"github.com/pkg/profile"
)

// Flags
var (
	ctx, cancel     = context.WithCancel(context.Background())
	cpuProfiling    = flag.Bool("cp", false, "if yes, cpu profiling is enabled")
	dataTypes       = astiflag.NewStringsMap()
	format          = flag.String("f", "", "the format")
	inputPath       = flag.String("i", "", "the input path")
	memoryProfiling = flag.Bool("mp", false, "if yes, memory profiling is enabled")
)

func main() {
	// Init
	flag.Var(dataTypes, "d", "the datatypes whitelist")
	var s = astiflag.Subcommand()
	flag.Parse()
	astilog.FlagInit()

	// Handle signals
	handleSignals()

	// Start profiling
	if *cpuProfiling {
		defer profile.Start(profile.CPUProfile, profile.ProfilePath(".")).Stop()
	} else if *memoryProfiling {
		defer profile.Start(profile.MemProfile, profile.ProfilePath(".")).Stop()
	}

	// Build the reader
	var r io.Reader
	var err error
	if r, err = buildReader(); err != nil {
		astilog.Error(errors.Wrap(err, "astits: parsing input failed"))
		return
	}

	// Make sure the reader is closed properly
	if c, ok := r.(io.Closer); ok {
		defer c.Close()
	}

	// Create the demuxer
	var dmx = astits.New(ctx, r)

	// Switch on subcommand
	switch s {
	case "data":
		// Fetch data
		if err = data(dmx); err != nil {
			astilog.Error(errors.Wrap(err, "astits: fetching data failed"))
			return
		}
	default:
		// Fetch the programs
		var pgms []*Program
		if pgms, err = programs(dmx); err != nil {
			astilog.Error(errors.Wrap(err, "astits: fetching programs failed"))
			return
		}

		// Print
		switch *format {
		case "json":
			var e = json.NewEncoder(os.Stdout)
			e.SetIndent("", "  ")
			if err = e.Encode(pgms); err != nil {
				astilog.Error(errors.Wrap(err, "astits: json encoding to stdout failed"))
				return
			}
		default:
			fmt.Println("Programs are:")
			for _, pgm := range pgms {
				fmt.Printf("* %s\n", pgm)
			}
		}
	}
}

func handleSignals() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch)
	go func() {
		for s := range ch {
			astilog.Debugf("Received signal %s", s)
			switch s {
			case syscall.SIGABRT, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM:
				cancel()
			}
			return
		}
	}()
}

func buildReader() (r io.Reader, err error) {
	// Validate input
	if len(*inputPath) <= 0 {
		err = errors.New("Use -i to indicate an input path")
		return
	}

	// Parse input
	var u *url.URL
	if u, err = url.Parse(*inputPath); err != nil {
		err = errors.Wrap(err, "astits: parsing input path failed")
		return
	}

	// Switch on scheme
	switch u.Scheme {
	case "udp":
		// Resolve addr
		var addr *net.UDPAddr
		if addr, err = net.ResolveUDPAddr("udp", u.Host); err != nil {
			err = errors.Wrapf(err, "astits: resolving udp addr %s failed", u.Host)
			return
		}

		// Listen to multicast UDP
		var c *net.UDPConn
		if c, err = net.ListenMulticastUDP("udp", nil, addr); err != nil {
			err = errors.Wrapf(err, "astits: listening on multicast udp addr %s failed", u.Host)
			return
		}
		c.SetReadBuffer(1000)
		r = c
	default:
		// Open file
		var f *os.File
		if f, err = os.Open(*inputPath); err != nil {
			err = errors.Wrapf(err, "astits: opening %s failed", *inputPath)
			return
		}
		r = f<commit_msg>Updated UDP reader<commit_after>package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/asticode/go-astilog"
	"github.com/asticode/go-astitools/flag"
	"github.com/asticode/go-astits"
	"github.com/asticode/go-astiudp"
	"github.com/pkg/errors"
	"github.com/pkg/profile"
)

// Flags
var (
	ctx, cancel     = context.WithCancel(context.Background())
	cpuProfiling    = flag.Bool("cp", false, "if yes, cpu profiling is enabled")
	dataTypes       = astiflag.NewStringsMap()
	format          = flag.String("f", "", "the format")
	inputPath       = flag.String("i", "", "the input path")
	memoryProfiling = flag.Bool("mp", false, "if yes, memory profiling is enabled")
)

func main() {
	// Init
	flag.Var(dataTypes, "d", "the datatypes whitelist")
	var s = astiflag.Subcommand()
	flag.Parse()
	astilog.FlagInit()

	// Handle signals
	handleSignals()

	// Start profiling
	if *cpuProfiling {
		defer profile.Start(profile.CPUProfile, profile.ProfilePath(".")).Stop()
	} else if *memoryProfiling {
		defer profile.Start(profile.MemProfile, profile.ProfilePath(".")).Stop()
	}

	// Build the reader
	var r io.Reader
	var err error
	if r, err = buildReader(ctx); err != nil {
		astilog.Error(errors.Wrap(err, "astits: parsing input failed"))
		return
	}

	// Make sure the reader is closed properly
	if c, ok := r.(io.Closer); ok {
		defer c.Close()
	}

	// Create the demuxer
	var dmx = astits.New(ctx, r)

	// Switch on subcommand
	switch s {
	case "data":
		// Fetch data
		if err = data(dmx); err != nil {
			astilog.Error(errors.Wrap(err, "astits: fetching data failed"))
			return
		}
	default:
		// Fetch the programs
		var pgms []*Program
		if pgms, err = programs(dmx); err != nil {
			astilog.Error(errors.Wrap(err, "astits: fetching programs failed"))
			return
		}

		// Print
		switch *format {
		case "json":
			var e = json.NewEncoder(os.Stdout)
			e.SetIndent("", "  ")
			if err = e.Encode(pgms); err != nil {
				astilog.Error(errors.Wrap(err, "astits: json encoding to stdout failed"))
				return
			}
		default:
			fmt.Println("Programs are:")
			for _, pgm := range pgms {
				fmt.Printf("* %s\n", pgm)
			}
		}
	}
}

func handleSignals() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch)
	go func() {
		for s := range ch {
			astilog.Debugf("Received signal %s", s)
			switch s {
			case syscall.SIGABRT, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM:
				cancel()
			}
			return
		}
	}()
}

func buildReader(ctx context.Context) (r io.Reader, err error) {
	// Validate input
	if len(*inputPath) <= 0 {
		err = errors.New("Use -i to indicate an input path")
		return
	}

	// Parse input
	var u *url.URL
	if u, err = url.Parse(*inputPath); err != nil {
		err = errors.Wrap(err, "astits: parsing input path failed")
		return
	}

	// Switch on scheme
	switch u.Scheme {
	case "udp":
		// Resolve addr
		var addr *net.UDPAddr
		if addr, err = net.ResolveUDPAddr("udp", u.Host); err != nil {
			err = errors.Wrapf(err, "astits: resolving udp addr %s failed", u.Host)
			return
		}

		// Listen to multicast UDP
		var c *net.UDPConn
		if c, err = net.ListenMulticastUDP("udp", nil, addr); err != nil {
			err = errors.Wrapf(err, "astits: listening on multicast udp addr %s failed", u.Host)
			return
		}
		c.SetReadBuffer(4096)

		// Initialize UDP reader
		// It will read 4096 bytes at each iteration, and will store up to 2MB in its buffer
		var mr = astiudp.NewReader(ctx, c, 4096, 2048*1024)

		// Pipe reader
		go mr.Pipe()
		r = mr
	default:
		// Open file
		var f *os.File
		if f, err = os.Open(*inputPath); err != nil {
			err = errors.Wrapf(err, "astits: opening %s failed", *inputPath)
			return
		}
		r = f
{'fraction_non_alphanumeric': 0.15090403337969402, 'fraction_numerical': 0.004033379694019471, 'mean_word_length': 3.135135135135135, 'pattern_counts': {'":': 6, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5472291', 'n_tokens_mistral': 2958, 'n_tokens_neox': 2673, 'n_words': 717}