68 lines
1.5 KiB
Go
68 lines
1.5 KiB
Go
package prober
|
|
|
|
import (
|
|
"bufio"
|
|
"dtool/utils"
|
|
"fmt"
|
|
"os"
|
|
"sync"
|
|
)
|
|
|
|
//type ProbeTask func(string) interface{}
|
|
|
|
func output_process(output chan interface{}, file string, wg *sync.WaitGroup) {
|
|
f, err := os.Create(file)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
defer f.Close()
|
|
writer := bufio.NewWriter(f)
|
|
for {
|
|
if data, ok := <-output; ok {
|
|
str, err := OutputHandler(data)
|
|
if err != nil {
|
|
fmt.Printf("Error generating output: %s\n", err)
|
|
continue
|
|
}
|
|
_, err = writer.WriteString(str + "\n")
|
|
if err != nil {
|
|
fmt.Printf("Error writing file: %s\n", err)
|
|
}
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
writer.Flush()
|
|
wg.Done()
|
|
}
|
|
|
|
func concurrent_execution[T any](fn func(string, string) T, domain string, input chan string, output chan interface{}, wg *sync.WaitGroup) {
|
|
for {
|
|
if ip, ok := <-input; ok {
|
|
data := fn(ip, domain)
|
|
output <- data
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
wg.Done()
|
|
}
|
|
|
|
func CreateTask[T any](fn func(string, string) T, domain string, input_file string, output_file string, concurrent_num int) {
|
|
input_pool := make(chan string, 500)
|
|
output_pool := make(chan interface{}, 500)
|
|
var probe_tasks sync.WaitGroup
|
|
var store_tasks sync.WaitGroup
|
|
|
|
go utils.RetrieveLines(input_pool, input_file)
|
|
probe_tasks.Add(concurrent_num)
|
|
for i := 0; i < concurrent_num; i++ {
|
|
go concurrent_execution(fn, domain, input_pool, output_pool, &probe_tasks)
|
|
}
|
|
store_tasks.Add(1)
|
|
go output_process(output_pool, output_file, &store_tasks)
|
|
probe_tasks.Wait()
|
|
close(output_pool)
|
|
store_tasks.Wait()
|
|
}
|