import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from wordcloud import WordCloud
from dotenv import load_dotenv
import os

# Load environment variables from .env file
load_dotenv()
access_token = os.getenv("HUGGINGFACE_ACCESS_TOKEN")

# Page configuration
st.set_page_config(page_title="ReconNinja Wordlists", page_icon="💬", layout="wide")

# Sidebar for navigation
def display_sidebar():
    st.sidebar.title("Navigation")
    options = ["Wordlist Generator", "Statistics", "Security Analysis"]
    choice = st.sidebar.radio("Go to", options)
    return choice

# Header section
def display_header():
    st.title("💬 ReconNinja Wordlists")
    st.subheader("Tailored wordlists for efficient penetration testing")
    st.markdown("""
        This application generates customized wordlists for use in network reconnaissance and penetration testing. 
        Adjust the parameters to generate wordlists suited for your specific testing scenario.
    """)

# Sidebar for user input
def get_user_inputs():
    st.sidebar.header("Customize Your Wordlist")
    st.sidebar.markdown("""
        Adjust the following parameters to create wordlists optimized for your penetration testing tasks.
    """)
    wordlist_size = st.sidebar.slider("Wordlist Size", min_value=50, max_value=10000, value=1000, step=50)
    min_length = st.sidebar.slider("Minimum Word Length", min_value=3, max_value=12, value=6)
    max_length = st.sidebar.slider("Maximum Word Length", min_value=3, max_value=12, value=8)
    include_special_chars = st.sidebar.checkbox("Include Special Characters", value=False)
    include_numbers = st.sidebar.checkbox("Include Numbers", value=True)
    
    return wordlist_size, min_length, max_length, include_special_chars, include_numbers

# Wordlist generation logic (mock-up for your project)
def generate_wordlist(size, min_length, max_length, special_chars=False, numbers=True):
    words = []
    for _ in range(size):
        word = ''.join(np.random.choice(list("abcdefghijklmnopqrstuvwxyz"), size=np.random.randint(min_length, max_length)))
        if special_chars:
            word += np.random.choice(["!", "@", "#", "$", "%"])
        if numbers:
            word += np.random.choice([str(i) for i in range(10)])
        words.append(word)
    return words

# Wordlist generation and display
def generate_and_display_wordlist(wordlist_size, min_length, max_length, include_special_chars, include_numbers):
    try:
        # Generate the wordlist
        wordlist = generate_wordlist(
            size=wordlist_size, 
            min_length=min_length, 
            max_length=max_length, 
            special_chars=include_special_chars, 
            numbers=include_numbers
        )
        
        # Display a preview of the wordlist
        st.write(f"Preview of {wordlist_size} words:")
        st.dataframe(pd.DataFrame(wordlist[:20], columns=["Generated Words"]))  # Display first 20 words

        # Provide a download link for the full wordlist
        st.markdown("### Download Full Wordlist")
        csv_data = pd.Series(wordlist).to_csv(index=False).encode()
        st.download_button(
            label="Download Wordlist as CSV",
            data=csv_data,
            file_name="reconninja_wordlist.csv",
            mime="text/csv"
        )

        return wordlist

    except Exception as e:
        st.error(f"Error generating wordlist: {e}")
        return None

# Visualizing the wordlist statistics
def display_wordlist_statistics(wordlist):
    if wordlist:
        st.header("Wordlist Statistics")
        
        # Calculate and display word length distribution
        word_lengths = [len(word) for word in wordlist]
        word_length_df = pd.DataFrame(word_lengths, columns=["Word Length"])

        # Bar Chart for Word Length Distribution
        st.subheader("Word Length Distribution")
        fig, ax = plt.subplots(figsize=(8, 6))
        sns.countplot(x=word_length_df["Word Length"], ax=ax, palette="viridis")
        ax.set_title("Frequency of Word Lengths")
        ax.set_xlabel("Word Length")
        ax.set_ylabel("Frequency")
        st.pyplot(fig)

        # Word Cloud of Words
        st.subheader("Word Cloud")
        wordcloud = WordCloud(width=800, height=400, background_color="white").generate(" ".join(wordlist))
        st.image(wordcloud.to_array(), use_column_width=True)

# Analyze wordlist security (entropy)
def analyze_wordlist_security(wordlist):
    if wordlist:
        st.header("Analyze Wordlist Security")
        
        entropy_slider = st.slider(
            "Select Entropy Multiplier", 
            min_value=1.0, 
            max_value=10.0, 
            value=3.0, 
            step=0.1
        )
        
        # Simulate password entropy calculation
        entropy = np.log2(len(wordlist) ** entropy_slider)
        st.write(f"Estimated Entropy: {entropy:.2f} bits")

        # Security analysis feedback
        if entropy < 50:
            st.warning("Low entropy detected! This wordlist might be vulnerable to brute-force attacks.")
        else:
            st.success("Good entropy! This wordlist is secure against most brute-force attempts.")

# Footer section
def display_footer():
    st.markdown("---")
    st.markdown(
        "Made with ❤️ by Canstralian. For more information on ReconNinja, visit our [GitHub](https://github.com/Canstralian)."
    )

# Main application function
def main():
    choice = display_sidebar()

    display_header()

    if 'wordlist' not in st.session_state:
        st.session_state.wordlist = None  # Initialize wordlist if it doesn't exist

    if choice == "Wordlist Generator":
        wordlist_size, min_length, max_length, include_special_chars, include_numbers = get_user_inputs()
        wordlist = generate_and_display_wordlist(
            wordlist_size, min_length, max_length, include_special_chars, include_numbers
        )
        # Store wordlist in session_state
        st.session_state.wordlist = wordlist

    elif choice == "Statistics":
        if st.session_state.wordlist is None:
            st.warning("Please generate a wordlist first!")
        else:
            display_wordlist_statistics(st.session_state.wordlist)

    elif choice == "Security Analysis":
        if st.session_state.wordlist is None:
            st.warning("Please generate a wordlist first!")
        else:
            analyze_wordlist_security(st.session_state.wordlist)

    display_footer()

if __name__ == "__main__":
    main()