如何在不退出程序的情況下檢查數字是否為質數，並再次詢問用戶是否不是質數？

Question

我正在嘗試檢查數字是否是質數，作為Diffie Hellman算法的一部分，如果不是質數，請再次要求用戶輸入另一個數字而不退出程序。

這是我的代碼：

def prime(p):
    for i in range(2,p):
        if (p % i) == 0:
            return(False)
            break
        else:
            return(True)


    def primRoots(p):
        roots = []
        required_set = set(num for num in range (1, p))

        for g in range(1, p):
            actual_set = set(pow(g, powers) % p for powers in range (1, p))
            if required_set == actual_set:
                roots.append(g)           
        return roots

    p=int(input("enter any prime no:"))
    check=prime(p)
    if(check==True):

        primitive_roots = primRoots(p)

        g=primitive_roots[0]
        print(g)
        x=int(input("Alice chooses value of X as:"))
        y=int(input("Bob chooses value of y as:"))
        r1=(g**x) % p
        r2=(g**y) % p

        print("value of r1 is",r1)
        print("value of r2 is",r2)

        a=x*y

        k1=(r2**x) % p
        print("k1 is",k1)
        k2=(r1**y)% p
        print("k2 is",k2)
        k=(g**a)%p
        print("shared key is",k)    
    else:
        print("It is not a prime num,enter again")
        prime(p)

Answer 1

user_input = int(input("Enter an integer !  "))

def prime(n):
    isPrime = True
    for num in range(2, int(n/2)):
        if n % num == 0:
            isPrime = False
    return isPrime


while (prime(user_input)== False):
    user_input = int(input("==> Your input is incorrect, please enter  an other integer !"))
print("{} is prime ".format(user_input))