This commit is contained in:
expvintl
2025-02-21 20:28:02 +08:00
commit 6e5b3875f9
7 changed files with 229 additions and 0 deletions

12
go.mod Normal file
View File

@ -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
)

12
go.sum Normal file
View File

@ -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=

92
main.go Normal file
View File

@ -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)
}

15
models/models_info.go Normal file
View File

@ -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"`
}

31
readme.md Normal file
View File

@ -0,0 +1,31 @@
<div align="center">
# Ollama Scanner (扫描器)
</div>
一个多线程的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)
```

40
utils/file.go Normal file
View File

@ -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])
}

27
utils/goroutine_pool.go Normal file
View File

@ -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)
}