93 lines
2.3 KiB
Go
93 lines
2.3 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"ollamaScan/models"
|
|
"ollamaScan/utils"
|
|
|
|
"github.com/schollz/progressbar/v3"
|
|
)
|
|
|
|
func ScanIP(ip string, port int) models.OllamaInfo {
|
|
c := &http.Client{
|
|
Timeout: time.Second * 2,
|
|
}
|
|
req, err := http.NewRequest("GET", "http://"+ip+":"+strconv.Itoa(port)+"/api/tags", nil)
|
|
if err != nil {
|
|
return models.OllamaInfo{}
|
|
}
|
|
r, err := c.Do(req)
|
|
if err != nil {
|
|
return models.OllamaInfo{}
|
|
}
|
|
defer r.Body.Close()
|
|
d, _ := io.ReadAll(r.Body)
|
|
info := models.OllamaInfo{
|
|
Host: ip,
|
|
Port: port,
|
|
}
|
|
err = json.Unmarshal(d, &info)
|
|
if err != nil {
|
|
return models.OllamaInfo{}
|
|
}
|
|
return info
|
|
}
|
|
func main() {
|
|
ipList := flag.String("l", "./ips.txt", "IP List File (default=./ips.txt)")
|
|
threads := flag.Int("t", 50, "Thread Num (default=50)")
|
|
outPath := flag.String("o", "./out.txt", "Output Save File (default=./out.txt)")
|
|
useJson := flag.Bool("json", false, "Output JSON Format (default=false)")
|
|
d, err := utils.ReadFile(*ipList)
|
|
if err != nil {
|
|
fmt.Println("Error:", err)
|
|
return
|
|
}
|
|
list := strings.Split(d, "\n")
|
|
pro := progressbar.Default(int64(len(list)), "Scanning...")
|
|
pool := utils.PoolInfo{}
|
|
pool.NewPool(*threads)
|
|
results := make([]models.OllamaInfo, 2000)
|
|
for i := 0; i < len(list); i++ {
|
|
pool.AddTask(func() {
|
|
info := ScanIP(list[i], 11434)
|
|
if info.Port != 0 {
|
|
results = append(results, info)
|
|
}
|
|
pro.Add(1)
|
|
})
|
|
}
|
|
defer pool.Pool.Release()
|
|
pool.TaskWaitGroup.Wait()
|
|
if *useJson {
|
|
buff := strings.Builder{}
|
|
out, _ := json.Marshal(results)
|
|
buff.Write(out)
|
|
if err := utils.WriteFile(*outPath, buff.String()); err != nil {
|
|
fmt.Println("Save Error:", err)
|
|
}
|
|
} else {
|
|
buff := strings.Builder{}
|
|
for _, i := range results {
|
|
if i.Port == 0 && len(i.Models) == 0 { //skip invalid data
|
|
continue
|
|
}
|
|
buff.WriteString("http://" + i.Host + ":" + strconv.Itoa(i.Port) + "/\n")
|
|
for _, j := range i.Models {
|
|
buff.WriteString("\t" + j.Name + " ModelSize:" + utils.FormatBytes(j.Size) + " " + " ParamSize:" + j.Details.ParameterSize + " QuantLevel:" + j.Details.QuantizationLevel + "\n")
|
|
}
|
|
}
|
|
if err := utils.WriteFile(*outPath, buff.String()); err != nil {
|
|
fmt.Println("Save Error:", err)
|
|
}
|
|
}
|
|
fmt.Println("Done! Saved to", outPath)
|
|
}
|