From e1923ee683cc58da147a219113ccd5e809cc3d07 Mon Sep 17 00:00:00 2001
From: TsMask <340112800@qq.com>
Date: Wed, 13 Sep 2023 14:35:10 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20UDM=E5=AF=BC=E5=85=A5=E6=96=87=E4=BB=B6?=
 =?UTF-8?q?=E6=88=90=E5=8A=9F=E7=9B=B8=E5=BA=94=E5=8E=9F=E5=A7=8B=E6=96=87?=
 =?UTF-8?q?=E6=9C=AC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 features/udm_user/api_udm_user.go             | 26 ++++++++++++-------
 .../udm_user/service/service_udm_auth_user.go |  2 +-
 .../udm_user/service/service_udm_sub_user.go  |  2 +-
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/features/udm_user/api_udm_user.go b/features/udm_user/api_udm_user.go
index 676e6471..c8af06c1 100644
--- a/features/udm_user/api_udm_user.go
+++ b/features/udm_user/api_udm_user.go
@@ -513,14 +513,18 @@ func (s *UdmUserApi) UdmAuthUserImport(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	neInfo, err := NeInfoByUDM(neId)
+	// 获取文件名
+	_, fileHeader, err := r.FormFile("file")
 	if err != nil {
 		ctx.JSON(w, 200, result.ErrMsg(err.Error()))
 		return
 	}
+	if !strings.HasSuffix(fileHeader.Filename, ".csv") {
+		ctx.JSON(w, 200, result.ErrMsg("请上传CSV格式文件，内容字段imsi, ki, algo, amf, opc"))
+		return
+	}
 
-	// 获取文件名
-	_, fileHeader, err := r.FormFile("file")
+	neInfo, err := NeInfoByUDM(neId)
 	if err != nil {
 		ctx.JSON(w, 200, result.ErrMsg(err.Error()))
 		return
@@ -558,8 +562,7 @@ func (s *UdmUserApi) UdmAuthUserImport(w http.ResponseWriter, r *http.Request) {
 		neId = "-"
 		s.authUser.InsertCSV(neId, data)
 	}
-	// ctx.JSON(w, 200, result.OkData(data))
-	ctx.FileAttachment(w, r, filePath, fileName)
+	ctx.JSON(w, 200, result.OkData(data))
 }
 
 // UDM签约用户
@@ -1038,14 +1041,18 @@ func (s *UdmUserApi) UdmSubUserImport(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	neInfo, err := NeInfoByUDM(neId)
+	// 获取文件名
+	_, fileHeader, err := r.FormFile("file")
 	if err != nil {
 		ctx.JSON(w, 200, result.ErrMsg(err.Error()))
 		return
 	}
+	if !strings.HasSuffix(fileHeader.Filename, ".csv") {
+		ctx.JSON(w, 200, result.ErrMsg("请上传CSV格式文件，内容字段imsi, msisdn, ambr, nssai, arfb, sar,rat, cn, smf_sel, sm_dat,eps_dat"))
+		return
+	}
 
-	// 获取文件名
-	_, fileHeader, err := r.FormFile("file")
+	neInfo, err := NeInfoByUDM(neId)
 	if err != nil {
 		ctx.JSON(w, 200, result.ErrMsg(err.Error()))
 		return
@@ -1083,6 +1090,5 @@ func (s *UdmUserApi) UdmSubUserImport(w http.ResponseWriter, r *http.Request) {
 		neId = "-"
 		s.subUser.InsertCSV(neId, data)
 	}
-	// ctx.JSON(w, 200, result.OkData(data))
-	ctx.FileAttachment(w, r, filePath, fileName)
+	ctx.JSON(w, 200, result.OkData(data))
 }
diff --git a/features/udm_user/service/service_udm_auth_user.go b/features/udm_user/service/service_udm_auth_user.go
index 82b5de23..c6651ba8 100644
--- a/features/udm_user/service/service_udm_auth_user.go
+++ b/features/udm_user/service/service_udm_auth_user.go
@@ -69,7 +69,7 @@ func (r *ServiceUdmAuthUser) InsertCSV(neID string, data []map[string]string) in
 		if s, ok := v["opc"]; ok {
 			authUser.Opc = s
 		}
-		r.repoAuthUser.Insert(neID, authUser)
+		num += r.repoAuthUser.Insert(neID, authUser)
 	}
 	return num
 }
diff --git a/features/udm_user/service/service_udm_sub_user.go b/features/udm_user/service/service_udm_sub_user.go
index 2e8839a2..f265eacf 100644
--- a/features/udm_user/service/service_udm_sub_user.go
+++ b/features/udm_user/service/service_udm_sub_user.go
@@ -86,7 +86,7 @@ func (r *ServiceUdmSubUser) InsertCSV(neID string, data []map[string]string) int
 		if s, ok := v["eps_dat"]; ok {
 			subUser.EpsDat = s
 		}
-		r.repoSunUser.Insert(neID, subUser)
+		num += r.repoSunUser.Insert(neID, subUser)
 	}
 	return num
 }