263 lines
6.5 KiB
Go
263 lines
6.5 KiB
Go
package laosiji
|
|
|
|
import (
|
|
"91porn-server/common/log"
|
|
"91porn-server/web/webg"
|
|
"bytes"
|
|
|
|
"context"
|
|
"crypto/aes"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io/ioutil"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
Name string
|
|
Appid string
|
|
APIKey string
|
|
APIUrl string
|
|
IMAGEYUAN string
|
|
NoticeURL string
|
|
)
|
|
|
|
// Config 是老司机接口配置。所有凭证必须由各服务自己的配置文件注入。
|
|
type Config struct {
|
|
AppID string
|
|
APIKey string
|
|
APIUrl string
|
|
ImageYuan string
|
|
NoticeURL string
|
|
}
|
|
|
|
func Init(lsjCfg webg.GlobalConfig) {
|
|
InitConfig(Config{
|
|
AppID: lsjCfg.LSJ.AppID,
|
|
APIKey: lsjCfg.LSJ.APIKey,
|
|
APIUrl: lsjCfg.LSJ.APIUrl,
|
|
ImageYuan: lsjCfg.LSJ.ImageYuan,
|
|
})
|
|
}
|
|
|
|
// InitConfig 供不依赖 web 配置结构的服务（例如 skd）初始化老司机客户端。
|
|
func InitConfig(cfg Config) {
|
|
Name = "老司机"
|
|
Appid = strings.TrimSpace(cfg.AppID)
|
|
APIKey = strings.TrimSpace(cfg.APIKey)
|
|
APIUrl = strings.TrimRight(strings.TrimSpace(cfg.APIUrl), "/")
|
|
IMAGEYUAN = strings.TrimRight(strings.TrimSpace(cfg.ImageYuan), "/")
|
|
NoticeURL = strings.TrimSpace(cfg.NoticeURL)
|
|
}
|
|
|
|
// Configured 判断调用老司机接口所需的配置是否完整。
|
|
func Configured() bool {
|
|
return Appid != "" && APIKey != "" && APIUrl != ""
|
|
}
|
|
|
|
func ensureConfigured() error {
|
|
if Configured() {
|
|
return nil
|
|
}
|
|
missing := make([]string, 0, 3)
|
|
if Appid == "" {
|
|
missing = append(missing, "appId")
|
|
}
|
|
if APIKey == "" {
|
|
missing = append(missing, "apiKey")
|
|
}
|
|
if APIUrl == "" {
|
|
missing = append(missing, "apiUrl")
|
|
}
|
|
return fmt.Errorf("laosiji configuration missing: %s", strings.Join(missing, ","))
|
|
}
|
|
|
|
type Response struct {
|
|
Status string `json:"status"`
|
|
Data string `json:"data"` // 如果成功，返回的这个数据是加密的
|
|
Time string `json:"time"`
|
|
Error string `json:"error"`
|
|
ErrorCode int `json:"errorCode"`
|
|
}
|
|
|
|
func QueryUndress(ctx context.Context, taskID string) (resp QueryUndressResponse, err error) {
|
|
endpoint := getQueryUrl()
|
|
err = post(ctx, endpoint, map[string]interface{}{
|
|
"task_id": taskID,
|
|
}, &resp)
|
|
if err != nil {
|
|
log.Error("GenerateUndress post fail", log.Any("taskID", taskID), log.E(err))
|
|
}
|
|
return
|
|
}
|
|
|
|
func GenerateUndress(ctx context.Context, req map[string]interface{}) (resp GenerateUndressResponse, err error) {
|
|
endpoint := getAiUndressGenerateUrl()
|
|
err = post(ctx, endpoint, req, &resp)
|
|
if err != nil {
|
|
log.Error("GenerateUndress post fail", log.Any("req", req), log.E(err))
|
|
}
|
|
return
|
|
}
|
|
|
|
func QueryTextToImage(ctx context.Context, taskID string) (resp QueryTextToImageResponse, err error) {
|
|
endpoint := getQueryUrl()
|
|
err = post(ctx, endpoint, map[string]interface{}{
|
|
"task_id": taskID,
|
|
}, &resp)
|
|
if err != nil {
|
|
log.Error("QueryTextToImage post fail", log.Any("taskID", taskID), log.E(err))
|
|
}
|
|
return
|
|
}
|
|
|
|
func GenerateTextToImage(ctx context.Context, req map[string]interface{}) (resp GenerateTextToImageResponse, err error) {
|
|
endpoint := getTextToImageGenerateUrl()
|
|
err = post(ctx, endpoint, req, &resp)
|
|
if err != nil {
|
|
log.Error("GenerateTextToImage post fail", log.Any("req", req), log.E(err))
|
|
}
|
|
return
|
|
}
|
|
|
|
func getQueryUrl() string {
|
|
return fmt.Sprintf("%s/lsjapi/ai/detail", APIUrl)
|
|
}
|
|
|
|
func getTextToImageGenerateUrl() string {
|
|
return fmt.Sprintf("%s/lsjapi/ai/generate", APIUrl)
|
|
}
|
|
|
|
func getAiUndressGenerateUrl() string {
|
|
return fmt.Sprintf("%s/lsjapi/ai/undress", APIUrl)
|
|
}
|
|
|
|
func post(ctx context.Context, endpoint string, req map[string]interface{}, response interface{}) (err error) {
|
|
if err = ensureConfigured(); err != nil {
|
|
return err
|
|
}
|
|
jsonData, err := json.Marshal(req)
|
|
if err != nil {
|
|
log.Warn("GenerateAiUndress", log.E(err))
|
|
return
|
|
}
|
|
encryptedData, err := encryptBase64(string(jsonData), APIKey)
|
|
if err != nil {
|
|
log.Warn("GenerateAiUndress encryptBase64", log.E(err))
|
|
return
|
|
}
|
|
// 创建 POST 请求
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewBuffer([]byte(encryptedData)))
|
|
if err != nil {
|
|
return
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
httpReq.Header.Set("Content-Length", fmt.Sprintf("%d", len(encryptedData)))
|
|
httpReq.Header.Set("appid", Appid)
|
|
|
|
// 发送请求
|
|
client := &http.Client{Timeout: 10 * time.Second}
|
|
resp, err := client.Do(httpReq)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// 读取响应体
|
|
body, err := ioutil.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
// 检查 HTTP 状态码
|
|
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
|
return fmt.Errorf("laosiji api http status:%d body:%s", resp.StatusCode, truncateToolBody(body, APIKey, Appid))
|
|
}
|
|
|
|
// 解析 JSON
|
|
var res Response
|
|
if err = json.Unmarshal(body, &res); err != nil {
|
|
return
|
|
}
|
|
if res.Status != "y" {
|
|
err = fmt.Errorf("errorCode:%v error:%s", res.ErrorCode, truncateToolText(res.Error, APIKey, Appid))
|
|
return
|
|
}
|
|
|
|
// 进行解密
|
|
dataStr, err := decryptBase64(res.Data, APIKey)
|
|
if err != nil {
|
|
return
|
|
}
|
|
err = json.Unmarshal([]byte(dataStr), response)
|
|
if err != nil {
|
|
return
|
|
}
|
|
return
|
|
}
|
|
|
|
// AES-128-ECB encrypt
|
|
func encryptBase64(input, key string) (string, error) {
|
|
if len(key) > 16 {
|
|
key = key[:16]
|
|
}
|
|
plainText := []byte(input)
|
|
keyBytes := []byte(key)
|
|
block, err := aes.NewCipher(keyBytes)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
blockSize := block.BlockSize()
|
|
plainText = pkcs7Padding(plainText, blockSize)
|
|
|
|
encrypted := make([]byte, len(plainText))
|
|
for bs, be := 0, blockSize; bs < len(plainText); bs, be = bs+blockSize, be+blockSize {
|
|
block.Encrypt(encrypted[bs:be], plainText[bs:be])
|
|
}
|
|
|
|
return base64.StdEncoding.EncodeToString(encrypted), nil
|
|
}
|
|
|
|
// AES-128-ECB decrypt
|
|
func decryptBase64(cipherText, key string) (string, error) {
|
|
if len(key) > 16 {
|
|
key = key[:16]
|
|
}
|
|
data, err := base64.StdEncoding.DecodeString(cipherText)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
block, err := aes.NewCipher([]byte(key))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
blockSize := block.BlockSize()
|
|
decrypted := make([]byte, len(data))
|
|
for bs, be := 0, blockSize; bs < len(data); bs, be = bs+blockSize, be+blockSize {
|
|
block.Decrypt(decrypted[bs:be], data[bs:be])
|
|
}
|
|
|
|
decrypted = pkcs7UnPadding(decrypted)
|
|
return string(decrypted), nil
|
|
}
|
|
|
|
// PKCS7Padding pads plaintext for AES ECB
|
|
func pkcs7Padding(src []byte, blockSize int) []byte {
|
|
padding := blockSize - len(src)%blockSize
|
|
padText := bytes.Repeat([]byte{byte(padding)}, padding)
|
|
return append(src, padText...)
|
|
}
|
|
|
|
// PKCS7UnPadding removes padding
|
|
func pkcs7UnPadding(src []byte) []byte {
|
|
length := len(src)
|
|
unpadding := int(src[length-1])
|
|
return src[:(length - unpadding)]
|
|
}
|