commit 6e5b3875f97dd89dd4cc5dcd59f5bbeb30bc2b23 Author: expvintl Date: Fri Feb 21 20:28:02 2025 +0800 init diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..156a813 --- /dev/null +++ b/go.mod @@ -0,0 +1,12 @@ +module ollamaScan + +go 1.22.5 + +require ( + github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect + github.com/panjf2000/ants v1.3.0 // indirect + github.com/rivo/uniseg v0.4.7 // indirect + github.com/schollz/progressbar/v3 v3.18.0 // indirect + golang.org/x/sys v0.30.0 // indirect + golang.org/x/term v0.29.0 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..e1477c5 --- /dev/null +++ b/go.sum @@ -0,0 +1,12 @@ +github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ= +github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw= +github.com/panjf2000/ants v1.3.0 h1:8pQ+8leaLc9lys2viEEr8md0U4RN6uOSUCE9bOYjQ9M= +github.com/panjf2000/ants v1.3.0/go.mod h1:AaACblRPzq35m1g3enqYcxspbbiOJJYaxU2wMpm1cXY= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/schollz/progressbar/v3 v3.18.0 h1:uXdoHABRFmNIjUfte/Ex7WtuyVslrw2wVPQmCN62HpA= +github.com/schollz/progressbar/v3 v3.18.0/go.mod h1:IsO3lpbaGuzh8zIMzgY3+J8l4C8GjO0Y9S69eFvNsec= +golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU= +golang.org/x/term v0.29.0/go.mod h1:6bl4lRlvVuDgSf3179VpIxBF0o10JUpXWOnI7nErv7s= diff --git a/main.go b/main.go new file mode 100644 index 0000000..e32a90f --- /dev/null +++ b/main.go @@ -0,0 +1,92 @@ +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) +} diff --git a/models/models_info.go b/models/models_info.go new file mode 100644 index 0000000..cf213aa --- /dev/null +++ b/models/models_info.go @@ -0,0 +1,15 @@ +package models + +type ModelInfo struct { + Name string `json:"name"` + Size uint64 `json:"size"` + Details struct { + ParameterSize string `json:"parameter_size"` + QuantizationLevel string `json:"quantization_level"` + } `json:"details"` +} +type OllamaInfo struct { + Host string `json:"host"` + Port int `json:"port"` + Models []ModelInfo `json:"models"` +} diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..5ea7e91 --- /dev/null +++ b/readme.md @@ -0,0 +1,31 @@ +
+ +# Ollama Scanner (扫描器) + +
+ +一个多线程的Ollama扫描器工具 +使用ants库实现的简单扫描器 + +## Build (构建) + +```text + go build . +``` + +## How to use? + +```text + ollamaScan + -t Num Threads (Default:50) + -l IP List File(Single IP Line) (Default:ips.txt) + -o Output File (Default: out.txt) + -json Save as Json Format (Default: False) +``` +```text + ollamaScan + -t 线程数 (默认:50) + -l IP列表文件(一行一个) (默认:ips.txt) + -o 输出文件 (默认: out.txt) + -json 保存为Json (默认: False) +``` \ No newline at end of file diff --git a/utils/file.go b/utils/file.go new file mode 100644 index 0000000..d306b42 --- /dev/null +++ b/utils/file.go @@ -0,0 +1,40 @@ +package utils + +import ( + "fmt" + "io" + "os" +) + +func ReadFile(path string) (string, error) { + file, err := os.Open(path) + if err != nil { + return "", nil + } + defer file.Close() + d, _ := io.ReadAll(file) + return string(d), nil +} + +func WriteFile(path string, content string) error { + file, err := os.Create(path) + if err != nil { + return err + } + defer file.Close() + file.WriteString(content) + return nil +} + +func FormatBytes(bytes uint64) string { + const unit = 1024 + if bytes < unit { + return fmt.Sprintf("%d B", bytes) + } + div, exp := uint64(unit), 0 + for n := bytes / unit; n >= unit; n /= unit { + div *= unit + exp++ + } + return fmt.Sprintf("%.1f %ciB", float64(bytes)/float64(div), "KMGTPE"[exp]) +} diff --git a/utils/goroutine_pool.go b/utils/goroutine_pool.go new file mode 100644 index 0000000..af2a87c --- /dev/null +++ b/utils/goroutine_pool.go @@ -0,0 +1,27 @@ +package utils + +import ( + "fmt" + "sync" + + "github.com/panjf2000/ants" +) + +type PoolInfo struct { + Pool *ants.Pool + MaxWorkers int + TaskWaitGroup sync.WaitGroup +} + +func (pool *PoolInfo) NewPool(num int) { + p, err := ants.NewPool(num) + if err != nil { + fmt.Println("Create Pool Error:", err) + return + } + pool.Pool = p + pool.MaxWorkers = num +} +func (pool *PoolInfo) AddTask(fun func()) { + pool.Pool.Submit(fun) +}