1
import os
2
import uuid
3
import json
4
import logging
5
import sys
6
from flask import Flask, request, redirect, render_template, session, url_for
7
from flask_limiter import Limiter
8
from flask_limiter.util import get_remote_address
9
import pymysql
10
from openai import OpenAI
11

12
openai_client = OpenAI()
13

14
app = Flask(__name__)
15
app.secret_key = os.getenv("SESSION_SECRET")
16

17
limiter = Limiter(
18
    get_remote_address,
19
    app=app,
20
    default_limits=[],
21
    storage_uri="redis://redis:6379",
22
)
23

24
# DB接続
25
def get_db():
26
    return pymysql.connect(
27
        host=os.getenv("MYSQL_HOST"),
28
        user=os.getenv("MYSQL_USER"),
29
        password=os.getenv("MYSQL_PASSWORD"),
30
        database=os.getenv("MYSQL_DATABASE"),
31
        charset='utf8mb4',
32
        autocommit=True,
33
        cursorclass=pymysql.cursors.DictCursor,
34
    )
35

36
def query_db(sql, args=(), fetchone=False):
37
    con = get_db()
38
    with con.cursor() as cur:
39
        cur.execute(sql, args)
40
        result = cur.fetchone() if fetchone else cur.fetchall()
41
    con.close()
42
    return result
43

44
def execute_db(sql, args=()):
45
    con = get_db()
46
    with con.cursor() as cur:
47
        cur.execute(sql, args)
48
    con.close()
49

50
@app.route('/')
51
def index():
52
    if 'user_id' in session:
53
        return redirect(f"/users/{session['user_id']}")
54
    return redirect(url_for('login'))
55

56
# ユーザー登録
57
@app.route('/register', methods=['GET', 'POST'])
58
def register():
59
    if request.method == 'POST':
60
        username = request.form['username']
61
        password = request.form['password']
62
        # username の重複をチェック
63
        existing = query_db("SELECT 1 FROM users WHERE username=%s", (username,), fetchone=True)
64
        if existing:
65
            return 'このユーザー名は既に使われています。', 409
66
        user_id = str(uuid.uuid4())
67
        execute_db("INSERT INTO users (id, username, password) VALUES (%s, %s, %s)", (user_id, username, password))
68
        session['user_id'] = user_id
69
        return redirect(f"/users/{user_id}")
70
    return render_template('register.html')
71

72
# ログイン
73
@app.route('/login', methods=['GET', 'POST'])
74
def login():
75
    if request.method == 'POST':
76
        username = request.form['username']
77
        password = request.form['password']
78
        user = query_db("SELECT * FROM users WHERE username=%s AND password=%s", (username, password), fetchone=True)
79
        if user:
80
            session['user_id'] = user['id']
81
            return redirect(f"/users/{user['id']}")
82
        return 'ユーザー名またはパスワードが間違っています。', 403
83
    return render_template('login.html')
84

85
# ログアウト
86
@app.route('/logout')
87
def logout():
88
    session.clear()
89
    return redirect(url_for('login'))
90

91
# ユーザーページ（自分のメモは非公開も表示、他人のメモは公開のみ）
92
@app.route('/users/<uid>')
93
def user_page(uid):
94
    current = session.get('user_id')
95
    if current == uid:
96
        sql = """
97
            SELECT id, body, visibility FROM memos WHERE user_id=%s AND visibility IN ('public','private')
98
            UNION
99
            SELECT id, '🔒秘密メモ' AS body, 'secret' AS visibility FROM memos WHERE user_id=%s AND visibility='secret'
100
        """
101
        memos = query_db(sql, (uid, uid))
102
    else:
103
        memos = query_db("SELECT id, body, visibility FROM memos WHERE user_id=%s AND visibility='public'", (uid,))
104
    return render_template('index.html', memos=memos)
105

106
# メモの詳細表示（secret の場合はパスワードを要求）
107
@app.route('/memo/<mid>', methods=['GET', 'POST'])
108
def memo_detail(mid):
109
    uid = session.get('user_id')
110
    memo = query_db('SELECT * FROM memos WHERE id=%s', (mid,), fetchone=True)
111
    if not memo:
112
        return 'Not found', 404
113
    if memo['user_id'] != uid:
114
        return 'Forbidden', 403
115
    if memo['visibility'] == 'secret':
116
        if request.method == 'POST' and request.form.get('password') == memo.get('password'):
117
            return render_template('detail.html', memo=memo, authorized=True)
118
        return render_template('detail.html', memo=memo, authorized=False) if request.method == 'GET' else ('Wrong password', 403)
119
    return render_template('detail.html', memo=memo, authorized=True)
120

121
# メモの作成
122
@app.route('/memo/create', methods=['GET', 'POST'])
123
def memo_create():
124
    uid = session.get('user_id')
125
    if not uid:
126
        return redirect('/')
127
    if request.method == 'POST':
128
        memo_count = query_db("SELECT COUNT(*) AS count FROM memos WHERE user_id=%s", (uid,), fetchone=True)['count']
129
        if memo_count >= 3:
130
            return "メモは3つまでしか作成できません。", 403
131

132
        body = request.form.get('body', '')
133
        if len(body) > 100:
134
            return "メモは100文字以下で入力してください。", 400
135

136
        visibility = request.form.get('visibility', 'public')
137
        password = request.form.get('password', '') if visibility == 'secret' else None
138
        mid = str(uuid.uuid4())
139
        execute_db(
140
            'INSERT INTO memos (id,user_id,body,visibility,password) VALUES (%s,%s,%s,%s,%s)',
141
            (mid, uid, body, visibility, password)
142
        )
143
        return redirect(f'/memo/{mid}')
144
    return render_template('create.html')
145

146
# 指定ユーザーのメモをキーワードで検索
147
def search_memos(keyword: str, include_secret: bool, user_id: str) -> list:
148
    visibilities = ("public","private","secret") if include_secret else ("public","private")
149
    placeholders = ','.join(['%s'] * len(visibilities))
150
    sql = f"SELECT id, body FROM memos WHERE user_id=%s AND visibility IN ({placeholders})"
151
    rows = query_db(sql, (user_id, *visibilities))
152
    return [r for r in rows if keyword.lower() in r['body'].lower()]
153

154
# 指定キーワードを含むメモの投稿者を取得
155
def get_author_by_body(keyword: str) -> list:
156
    row = query_db("SELECT user_id FROM memos WHERE body LIKE %s ORDER BY created_at ASC LIMIT 1", (f"%{keyword}%",), fetchone=True)
157
    return [{'user_id': row['user_id']}] if row else []
158

159
# RAG機能：検索や投稿者取得をfunction callingで実施
160
def rag(query: str, user_id: str) -> list:
161
    tools = [
162
        {
163
            'type': 'function',
164
            'function': {
165
                'name': 'search_memos',
166
                'description': 'Search for memos by keyword and visibility settings.',
167
                'parameters': {
168
                    'type': 'object',
169
                    'properties': {
170
                        'keyword': {'type': 'string'},
171
                        'include_secret': {'type': 'boolean'},
172
                        'target_uid': {'type': 'string'}
173
                    },
174
                    'required': ['keyword', 'include_secret', 'target_uid'],
175
                }
176
            }
177
        },
178
        {
179
            'type': 'function',
180
            'function': {
181
                'name': 'get_author_by_body',
182
                'description': 'Find the user who wrote a memo containing a given keyword.',
183
                'parameters': {
184
                    'type': 'object',
185
                    'properties': {
186
                        'keyword': {'type': 'string'}
187
                    },
188
                    'required': ['keyword']
189
                }
190
            }
191
        }
192
    ]
193
    response = openai_client.chat.completions.create(
194
        model='gpt-4o-mini',
195
        messages=[
196
            {'role': 'system', 'content': 'You are an assistant that helps search user memos using the available tools.'},
197
            {'role': 'assistant', 'content': 'Target User ID: ' + user_id},
198
            {'role': 'user', 'content': query}
199
        ],
200
        tools=tools,
201
        tool_choice='required',
202
        max_tokens=100,
203
    )
204
    choice = response.choices[0]
205
    if choice.message.tool_calls:
206
        call = choice.message.tool_calls[0]
207
        name = call.function.name
208
        args = json.loads(call.function.arguments)
209
        if name == 'search_memos':
210
            return search_memos(args.get('keyword', ''), args.get('include_secret', False), args.get('target_uid', ''))
211
        elif name == 'get_author_by_body':
212
            return get_author_by_body(args['keyword'])
213
    return []
214

215
# メモを文脈にして質問に答える
216
def answer_with_context(query: str, memos: list) -> str:
217
    context_text = "\n---\n".join([m['body'] for m in memos])
218
    prompt = f"""Here are your memos. Answer the following question based on them:
219

220
{context_text}
221

222
Question: {query}
223
"""
224
    response = openai_client.chat.completions.create(
225
        model='gpt-4o-mini',
226
        messages=[
227
            {'role': 'system', 'content': 'You are an assistant that answers questions using the user\'s memos as context.'},
228
            {'role': 'user', 'content': prompt}
229
        ],
230
        max_tokens=100,
231
    )
232
    content = response.choices[0].message.content.strip()
233
    return content
234

235
# RAGによるメモ検索
236
@app.route('/memo/search', methods=['GET'])
237
def search_form():
238
    uid = session.get('user_id')
239
    if not uid:
240
        return redirect('/')
241
    return render_template('search.html', answer=None, query='')
242

243
@app.route('/memo/search', methods=['POST'])
244
@limiter.limit("5 per minute")
245
def search():
246
    uid = session.get('user_id')
247
    if not uid:
248
        return redirect('/')
249
    query = request.form.get('query', '')
250
    memos = rag(query, uid)
251
    if not (memos and isinstance(memos, list)):
252
        answer = "関連するメモが見つかりませんでした。"
253
    else:
254
        if 'user_id' in memos[0]:
255
            answer = f"User ID: {memos[0]['user_id']}"
256
        else:
257
            answer = answer_with_context(query, memos)
258
            # 回答にFLAGが含まれている場合は警告を表示
259
            if "ctf4b" in answer:
260
                answer = "FLAGのメモは取得できません。"
261
    return render_template('search.html', answer=answer, query=query)
262

263
# ログ出力の設定
264
logging.basicConfig(
265
    level=logging.INFO,
266
    format="%(asctime)s [%(levelname)s] %(message)s",
267
    handlers=[logging.StreamHandler(sys.stdout)]
268
)

1
CREATE DATABASE IF NOT EXISTS memodb
2
  DEFAULT CHARACTER SET utf8mb4
3
  COLLATE utf8mb4_unicode_ci;
4

5
USE memodb;
6

7
CREATE TABLE IF NOT EXISTS users (
8
  id VARCHAR(36) PRIMARY KEY,
9
  username VARCHAR(255) UNIQUE,
10
  password TEXT
11
) CHARACTER SET utf8mb4
12
  COLLATE utf8mb4_unicode_ci;
13

14
CREATE TABLE IF NOT EXISTS memos (
15
  id VARCHAR(36) PRIMARY KEY,
16
  user_id VARCHAR(36),
17
  body TEXT,
18
  visibility ENUM('public','private','secret') NOT NULL,
19
  password TEXT,
20
  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
21
) CHARACTER SET utf8mb4
22
  COLLATE utf8mb4_unicode_ci;
23

24
INSERT IGNORE INTO users (id, username, password) VALUES
25
('dummy_admin_id', 'admin', 'dummy_admin_pass');
26

27
INSERT IGNORE INTO memos (id, user_id, body, visibility, password) VALUES
28
('dummy_admin_memo_id', 'dummy_admin_id', 'ctf4b{dummy_flag}', 'secret', 'dummy_admin_memo_pass');

1
import express from 'express';
2
import sanitizeHtml from 'sanitize-html';
3
import { marked } from 'marked';
4
import parse from 'url-parse';
5
import fs from 'fs';
6
import path from 'path';
7
import crypto from 'crypto';
8

9
const app   = express();
10
const posts = new Map();
11
const FLAG  = fs.readFileSync('./flag.txt','utf8').trim();
12

13
const emojiMap = {
14
  'smile': '😊',
15
  'heart': '❤️',
16
  'fire': '🔥',
17
  'thumbsup': '👍',
18
  'thumbsdown': '👎',
19
  'check': '✅',
20
  'x': '❌',
21
  'party': '🎉',
22
  'eyes': '👀',
23
  'thinking': '🤔',
24
  'cry': '😢',
25
  'laugh': '😂'
26
};
27

28
app.use(express.urlencoded({ extended: false }));
29
app.use(express.json());
30

31
function renderTemplate(templatePath, data) {
32
  let template = fs.readFileSync(templatePath, 'utf8');
33

34
  for (const [key, value] of Object.entries(data)) {
35
    const placeholder = `{{${key.toUpperCase()}}}`;
36
    template = template.replace(new RegExp(placeholder, 'g'), value);
37
  }
38

39
  return template;
40
}
41

42
app.get('/flag', (req,res)=> {
43
  const clientIP = req.socket.remoteAddress;
44
  const isLocalhost = clientIP === '127.0.0.1' ||
45
                     clientIP?.startsWith('172.20.');
46

47
  if (!isLocalhost) {
48
    return res.status(403).json({ error: 'Access denied.' });
49
  }
50

51
  if (req.headers.cookie !== 'user=admin') {
52
    return res.status(403).json({ error: 'Admin access required.' });
53
  }
54

55
  res.type('text/plain').send(FLAG);
56
});
57

58
app.get('/', (_req, res) => {
59
  const html = renderTemplate('./templates/index.html', {});
60
  res.send(html);
61
});
62

63
app.post('/', (req,res)=>{
64
  const { title='', md='' } = req.body;
65

66
  marked.setOptions({
67
    breaks: true,
68
    gfm: false
69
  });
70

71
  let html = marked.parse(md);
72

73
  html = sanitizeHtml(html, {
74
    allowedTags: ['h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p', 'a', 'ul', 'ol', 'li', 'blockquote', 'code', 'pre', 'em', 'strong', 'br'],
75
    allowedAttributes: {
76
      'a': ['href']
77
    }
78
  });
79

80
  html = processEmojis(html);
81

82
  const id = crypto.randomUUID().slice(0,8);
83
  posts.set(id,{
84
    title: title.replace(/[<>]/g, ''),
85
    html: html
86
  });
87
  res.redirect(`/post/${id}`);
88
});
89

90
app.get('/post/:id', (req,res)=>{
91
  const post = posts.get(req.params.id);
92
  if(!post) return res.sendStatus(404);
93

94
  const html = renderTemplate('./templates/post.html', {
95
    title: post.title,
96
    content: post.html
97
  });
98
  res.send(html);
99
});
100

101
app.use('/static', express.static(path.join(process.cwd(),'static')));
102

103
app.get('/api/posts', (req, res) => {
104
  const postList = Array.from(posts.entries()).map(([id, post]) => ({
105
    id,
106
    title: post.title,
107
    url: `/post/${id}`
108
  }));
109
  res.json(postList);
110
});
111

112
function processEmojis(html) {
113
  return html.replace(/:((?:https?:\/\/[^:]+|[^:]+)):/g, (match, name) => {
114
    if (emojiMap[name]) {
115
      return emojiMap[name];
116
    }
117

118
    if (name.match(/^https?:\/\//)) {
119
      try {
120
        const urlObj = new URL(name);
121
        const baseUrl = urlObj.origin + urlObj.pathname;
122
        const parsed = parse(name);
123
        const fragment = parsed.hash || '';
124
        const imgUrl = baseUrl + fragment;
125

126
        return `<img src="${imgUrl}" style="height:1.2em;vertical-align:middle;">`;
127
      } catch (e) {
128
        return match;
129
      }
130
    }
131

132
    return match;
133
  });
134
}
135

136
app.listen(50000, '0.0.0.0', ()=>console.log('Server running on http://localhost:50000'));

1
import express, { Request, Response } from "express";
2
import session from "express-session";
3
import path from "path";
4
import { db } from "./database";
5

6
declare module "express-session" {
7
  interface SessionData {
8
    userId?: number;
9
    username?: string;
10
  }
11
}
12

13
const app = express();
14
const PORT = process.env.PORT || 3000;
15

16
app.use(express.json());
17
app.use(express.urlencoded({ extended: true }));
18
app.use(express.static(path.join(__dirname, "../public")));
19

20
app.use(
21
  session({
22
    secret:
23
      process.env.SESSION_SECRET || "your-secret-key-change-in-production",
24
    resave: false,
25
    saveUninitialized: false,
26
    cookie: { secure: false, maxAge: 24 * 60 * 60 * 1000 },
27
  })
28
);
29

30
app.post("/api/register", async (req: Request, res: Response) => {
31
  try {
32
    const { username, password } = req.body;
33
    if (!username || !password) {
34
      return res.status(400).json({ error: "Username and password required" });
35
    }
36

37
    const existingUser = await db.findUser(username);
38
    if (existingUser) {
39
      return res.status(400).json({ error: "Username already exists" });
40
    }
41

42
    const userId = await db.createUser(username, password);
43
    req.session.userId = userId;
44
    req.session.username = username;
45

46
    res.json({ success: true, message: "Registration successful" });
47
  } catch (error) {
48
    res.status(500).json({ error: "Registration failed" });
49
  }
50
});
51

52
app.post("/api/login", async (req: Request, res: Response) => {
53
  try {
54
    const { username, password } = req.body;
55
    if (!username || !password) {
56
      return res.status(400).json({ error: "Username and password required" });
57
    }
58

59
    const user = await db.findUser(username);
60
    if (!user || !db.validatePassword(password, user.password_hash)) {
61
      return res.status(401).json({ error: "Invalid credentials" });
62
    }
63

64
    req.session.userId = user.userid;
65
    req.session.username = user.username;
66

67
    res.json({ success: true, message: "Login successful" });
68
  } catch (error) {
69
    res.status(500).json({ error: "Login failed" });
70
  }
71
});
72

73
app.post("/api/logout", (req: Request, res: Response) => {
74
  req.session.destroy((err) => {
75
    if (err) {
76
      return res.status(500).json({ error: "Logout failed" });
77
    }
78
    res.json({ success: true, message: "Logout successful" });
79
  });
80
});
81

82
app.post("/api/reset-request", async (req: Request, res: Response) => {
83
  try {
84
    const { username } = req.body;
85

86
    if (!username) {
87
      return res.status(400).json({ error: "Username is required" });
88
    }
89

90
    const user = await db.findUser(username);
91
    if (!user) {
92
      return res.status(404).json({ error: "User not found" });
93
    }
94

95
    await db.generateResetToken(user.userid);
96

97
    // TODO: send email to admin
98
    res.json({
99
      success: true,
100
      message:
101
        "Reset token has been generated. Please contact the administrator for the token.",
102
    });
103
  } catch (error) {
104
    console.error("Error generating reset token:", error);
105
    res.status(500).json({ error: "Internal server error" });
106
  }
107
});
108

109
app.post("/api/reset-password", async (req: Request, res: Response) => {
110
  try {
111
    const { username, token, newPassword } = req.body;
112
    if (!username || !token || !newPassword) {
113
      return res
114
        .status(400)
115
        .json({ error: "Username, token, and new password are required" });
116
    }
117

118
    const isValid = await db.validateResetTokenByUsername(username, token);
119

120
    if (!isValid) {
121
      return res.status(400).json({ error: "Invalid token" });
122
    }
123

124
    // TODO: implement
125
    // await db.updatePasswordByUsername(username, newPassword);
126

127
    // TODO: remove this
128
    const user = await db.findUser(username);
129
    if (!user) {
130
      return res.status(401).json({ error: "Invalid username" });
131
    }
132
    req.session.userId = user.userid;
133
    req.session.username = user.username;
134

135
    res.json({
136
      success: true,
137
      message: `The function to update the password is not implemented, so I will set you the ${user.username}'s session`,
138
    });
139
  } catch (error) {
140
    console.error("Password reset error:", error);
141
    res.status(500).json({ error: "Reset failed" });
142
  }
143
});
144

145
app.get("/api/get_flag", (req: Request, res: Response) => {
146
  if (!req.session.userId) {
147
    return res.status(401).json({ error: "Not authenticated" });
148
  }
149

150
  if (req.session.username === "admin") {
151
    res.json({ flag: process.env.FLAG || "ctf4B{**REDACTED**}" });
152
  } else {
153
    res.json({ message: "Hello user! Only admin can see the flag." });
154
  }
155
});
156

157
app.get("/api/status", (req: Request, res: Response) => {
158
  if (req.session.userId) {
159
    res.json({
160
      authenticated: true,
161
      username: req.session.username,
162
      isAdmin: req.session.username === "admin",
163
    });
164
  } else {
165
    res.json({ authenticated: false });
166
  }
167
});
168

169
app.get("*", (req: Request, res: Response) => {
170
  res.sendFile(path.join(__dirname, "../public/index.html"));
171
});
172

173
app.listen(PORT, () => {
174
  console.log(`Server running on port ${PORT}`);
175
});

1
import mysql from "mysql2/promise";
2
import bcrypt from "bcryptjs";
3
import { v4 as uuidv4 } from "uuid";
4

5
export interface User {
6
  userid: number;
7
  username: string;
8
  password_hash: string;
9
  reset_token: string | null;
10
}
11

12
class Database {
13
  private pool: mysql.Pool;
14
  private initialized: Promise<void>;
15

16
  constructor() {
17
    this.pool = mysql.createPool({
18
      host: process.env.DB_HOST || "localhost",
19
      port: parseInt(process.env.DB_PORT || "3306"),
20
      user: process.env.DB_USER || "root",
21
      password: process.env.DB_PASSWORD || "rootpassword",
22
      database: process.env.DB_NAME || "login4b",
23
      waitForConnections: true,
24
      connectionLimit: 10,
25
      queueLimit: 0,
26
    });
27
    this.initialized = this.init();
28
  }
29

30
  private async init() {
31
    try {
32
      await this.pool.execute(`
33
        CREATE TABLE IF NOT EXISTS users (
34
          userid INT AUTO_INCREMENT PRIMARY KEY,
35
          username VARCHAR(255) UNIQUE NOT NULL,
36
          password_hash VARCHAR(255) NOT NULL,
37
          reset_token VARCHAR(255)
38
        )
39
      `);
40

41
      // Check if admin user exists
42
      const [rows] = (await this.pool.execute(
43
        "SELECT COUNT(*) as count FROM users WHERE username = ?",
44
        ["admin"]
45
      )) as [any[], mysql.FieldPacket[]];
46

47
      if (rows[0].count === 0) {
48
        const adminHash = bcrypt.hashSync(
49
          process.env.ADMIN_PASSWORD || "admin_pass",
50
          10
51
        );
52
        await this.pool.execute(
53
          "INSERT INTO users (username, password_hash) VALUES (?, ?)",
54
          ["admin", adminHash]
55
        );
56
      }
57
    } catch (error) {
58
      console.error("Database initialization error:", error);
59
    }
60
  }
61

62
  async createUser(username: string, password: string): Promise<number> {
63
    await this.initialized;
64
    const hashedPassword = bcrypt.hashSync(password, 10);
65
    const [result] = (await this.pool.execute(
66
      "INSERT INTO users (username, password_hash) VALUES (?, ?)",
67
      [username, hashedPassword]
68
    )) as [mysql.ResultSetHeader, mysql.FieldPacket[]];
69
    return result.insertId;
70
  }
71

72
  async findUser(username: string): Promise<User | null> {
73
    await this.initialized;
74
    const [rows] = (await this.pool.execute(
75
      "SELECT * FROM users WHERE username = ?",
76
      [username]
77
    )) as [User[], mysql.FieldPacket[]];
78
    return rows[0] || null;
79
  }
80

81
  validatePassword(password: string, hash: string): boolean {
82
    return bcrypt.compareSync(password, hash);
83
  }
84

85
  async generateResetToken(userid: number): Promise<string> {
86
    await this.initialized;
87
    const timestamp = Math.floor(Date.now() / 1000);
88
    const token = `${timestamp}_${uuidv4()}`;
89

90
    await this.pool.execute(
91
      "UPDATE users SET reset_token = ? WHERE userid = ?",
92
      [token, userid]
93
    );
94
    return token;
95
  }
96

97
  async validateResetTokenByUsername(
98
    username: string,
99
    token: string
100
  ): Promise<boolean> {
101
    await this.initialized;
102
    const [rows] = (await this.pool.execute(
103
      "SELECT COUNT(*) as count FROM users WHERE username = ? AND reset_token = ?",
104
      [username, token]
105
    )) as [any[], mysql.FieldPacket[]];
106
    return rows[0].count > 0;
107
  }
108
}
109

110
export const db = new Database();