github.com/metacubex/mihomo@v1.18.5/component/geodata/init.go (about)

     1  package geodata
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  	"os"
     9  	"time"
    10  
    11  	mihomoHttp "github.com/metacubex/mihomo/component/http"
    12  	"github.com/metacubex/mihomo/component/mmdb"
    13  	C "github.com/metacubex/mihomo/constant"
    14  	"github.com/metacubex/mihomo/log"
    15  )
    16  
    17  var (
    18  	initGeoSite bool
    19  	initGeoIP   int
    20  	initASN     bool
    21  )
    22  
    23  func InitGeoSite() error {
    24  	if _, err := os.Stat(C.Path.GeoSite()); os.IsNotExist(err) {
    25  		log.Infoln("Can't find GeoSite.dat, start download")
    26  		if err := downloadGeoSite(C.Path.GeoSite()); err != nil {
    27  			return fmt.Errorf("can't download GeoSite.dat: %s", err.Error())
    28  		}
    29  		log.Infoln("Download GeoSite.dat finish")
    30  		initGeoSite = false
    31  	}
    32  	if !initGeoSite {
    33  		if err := Verify(C.GeositeName); err != nil {
    34  			log.Warnln("GeoSite.dat invalid, remove and download: %s", err)
    35  			if err := os.Remove(C.Path.GeoSite()); err != nil {
    36  				return fmt.Errorf("can't remove invalid GeoSite.dat: %s", err.Error())
    37  			}
    38  			if err := downloadGeoSite(C.Path.GeoSite()); err != nil {
    39  				return fmt.Errorf("can't download GeoSite.dat: %s", err.Error())
    40  			}
    41  		}
    42  		initGeoSite = true
    43  	}
    44  	return nil
    45  }
    46  
    47  func downloadGeoSite(path string) (err error) {
    48  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
    49  	defer cancel()
    50  	resp, err := mihomoHttp.HttpRequest(ctx, C.GeoSiteUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
    51  	if err != nil {
    52  		return
    53  	}
    54  	defer resp.Body.Close()
    55  
    56  	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644)
    57  	if err != nil {
    58  		return err
    59  	}
    60  	defer f.Close()
    61  	_, err = io.Copy(f, resp.Body)
    62  
    63  	return err
    64  }
    65  
    66  func downloadGeoIP(path string) (err error) {
    67  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
    68  	defer cancel()
    69  	resp, err := mihomoHttp.HttpRequest(ctx, C.GeoIpUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
    70  	if err != nil {
    71  		return
    72  	}
    73  	defer resp.Body.Close()
    74  
    75  	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644)
    76  	if err != nil {
    77  		return err
    78  	}
    79  	defer f.Close()
    80  	_, err = io.Copy(f, resp.Body)
    81  
    82  	return err
    83  }
    84  
    85  func InitGeoIP() error {
    86  	if C.GeodataMode {
    87  		if _, err := os.Stat(C.Path.GeoIP()); os.IsNotExist(err) {
    88  			log.Infoln("Can't find GeoIP.dat, start download")
    89  			if err := downloadGeoIP(C.Path.GeoIP()); err != nil {
    90  				return fmt.Errorf("can't download GeoIP.dat: %s", err.Error())
    91  			}
    92  			log.Infoln("Download GeoIP.dat finish")
    93  			initGeoIP = 0
    94  		}
    95  
    96  		if initGeoIP != 1 {
    97  			if err := Verify(C.GeoipName); err != nil {
    98  				log.Warnln("GeoIP.dat invalid, remove and download: %s", err)
    99  				if err := os.Remove(C.Path.GeoIP()); err != nil {
   100  					return fmt.Errorf("can't remove invalid GeoIP.dat: %s", err.Error())
   101  				}
   102  				if err := downloadGeoIP(C.Path.GeoIP()); err != nil {
   103  					return fmt.Errorf("can't download GeoIP.dat: %s", err.Error())
   104  				}
   105  			}
   106  			initGeoIP = 1
   107  		}
   108  		return nil
   109  	}
   110  
   111  	if _, err := os.Stat(C.Path.MMDB()); os.IsNotExist(err) {
   112  		log.Infoln("Can't find MMDB, start download")
   113  		if err := mmdb.DownloadMMDB(C.Path.MMDB()); err != nil {
   114  			return fmt.Errorf("can't download MMDB: %s", err.Error())
   115  		}
   116  	}
   117  
   118  	if initGeoIP != 2 {
   119  		if !mmdb.Verify(C.Path.MMDB()) {
   120  			log.Warnln("MMDB invalid, remove and download")
   121  			if err := os.Remove(C.Path.MMDB()); err != nil {
   122  				return fmt.Errorf("can't remove invalid MMDB: %s", err.Error())
   123  			}
   124  			if err := mmdb.DownloadMMDB(C.Path.MMDB()); err != nil {
   125  				return fmt.Errorf("can't download MMDB: %s", err.Error())
   126  			}
   127  		}
   128  		initGeoIP = 2
   129  	}
   130  	return nil
   131  }
   132  
   133  func InitASN() error {
   134  	if _, err := os.Stat(C.Path.ASN()); os.IsNotExist(err) {
   135  		log.Infoln("Can't find ASN.mmdb, start download")
   136  		if err := mmdb.DownloadASN(C.Path.ASN()); err != nil {
   137  			return fmt.Errorf("can't download ASN.mmdb: %s", err.Error())
   138  		}
   139  		log.Infoln("Download ASN.mmdb finish")
   140  		initASN = false
   141  	}
   142  	if !initASN {
   143  		if !mmdb.Verify(C.Path.ASN()) {
   144  			log.Warnln("ASN invalid, remove and download")
   145  			if err := os.Remove(C.Path.ASN()); err != nil {
   146  				return fmt.Errorf("can't remove invalid ASN: %s", err.Error())
   147  			}
   148  			if err := mmdb.DownloadASN(C.Path.ASN()); err != nil {
   149  				return fmt.Errorf("can't download ASN: %s", err.Error())
   150  			}
   151  		}
   152  		initASN = true
   153  	}
   154  	return nil
   155  }