github.com/chainreactors/fingers@v1.2.1/resources/loader.go (about)

     1  package resources
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  	"net/http"
     7  	"os"
     8  	"strings"
     9  
    10  	"gopkg.in/yaml.v3"
    11  )
    12  
    13  // LoadResource loads content from file path or HTTP/HTTPS URL
    14  func LoadResource(path string) ([]byte, error) {
    15  	// Check if it's a local file
    16  	if _, err := os.Stat(path); err == nil {
    17  		return ioutil.ReadFile(path)
    18  	}
    19  
    20  	// Check if it's a URL
    21  	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
    22  		resp, err := http.Get(path)
    23  		if err != nil {
    24  			return nil, fmt.Errorf("failed to fetch from URL: %w", err)
    25  		}
    26  		defer resp.Body.Close()
    27  
    28  		if resp.StatusCode != http.StatusOK {
    29  			return nil, fmt.Errorf("bad status: %s", resp.Status)
    30  		}
    31  
    32  		return ioutil.ReadAll(resp.Body)
    33  	}
    34  
    35  	return nil, fmt.Errorf("invalid resource path: %s (not a file or URL)", path)
    36  }
    37  
    38  // LoadFingersFromYAML loads fingerprints from YAML format file or URL
    39  // This function is specifically for loading custom fingerprints in YAML format
    40  func LoadFingersFromYAML(path string) ([]byte, error) {
    41  	content, err := LoadResource(path)
    42  	if err != nil {
    43  		return nil, err
    44  	}
    45  
    46  	// Validate that it's valid YAML by attempting to unmarshal
    47  	var test interface{}
    48  	if err := yaml.Unmarshal(content, &test); err != nil {
    49  		return nil, fmt.Errorf("invalid YAML format: %w", err)
    50  	}
    51  
    52  	return content, nil
    53  }