Ngăn chặn các cuộc tấn công SQL Injection với Python

Cứ vài năm một lần, Dự án Bảo mật Ứng dụng Web Mở (OWASP) lại xếp hạng các rủi ro bảo mật ứng dụng web nghiêm trọng nhất. Kể từ báo cáo đầu tiên, rủi ro khi tiêm luôn được đặt lên hàng đầu. Trong số tất cả các kiểu tiêm, Chèn SQL là một trong những vectơ tấn công phổ biến nhất và được cho là nguy hiểm nhất. Vì Python là một trong những ngôn ngữ lập trình phổ biến nhất trên thế giới nên việc biết cách bảo vệ khỏi việc tiêm SQL trong Python là rất quan trọng.

Trong hướng dẫn này, bạn sẽ tìm hiểu:

Chèn SQL Python là gì là gì và làm thế nào để ngăn chặn nó
Cách soạn truy vấn với cả ký tự và số nhận dạng dưới dạng tham số
Cách thực thi các truy vấn một cách an toàn trong cơ sở dữ liệu

Hướng dẫn này phù hợp với người dùng của tất cả các công cụ cơ sở dữ liệu . Các ví dụ ở đây sử dụng PostgreSQL, nhưng kết quả có thể được sao chép trong các hệ thống quản lý cơ sở dữ liệu khác (chẳng hạn như SQLite, MySQL, Microsoft SQL Server, Oracle, v.v.).

Phần thưởng miễn phí: 5 Suy nghĩ khi Làm chủ Python, một khóa học miễn phí dành cho các nhà phát triển Python, cho bạn thấy lộ trình và tư duy mà bạn sẽ cần để nâng các kỹ năng Python của mình lên một cấp độ tiếp theo.

Hiểu về Python SQL Injection

Các cuộc tấn công SQL Injection là một lỗ hổng bảo mật phổ biến đến nỗi xkcd huyền thoại webcomic đã dành một truyện tranh cho nó:

Tạo và thực thi truy vấn SQL là một nhiệm vụ phổ biến. Tuy nhiên, các công ty trên khắp thế giới thường mắc phải những sai lầm khủng khiếp khi soạn các câu lệnh SQL. Mặc dù lớp ORM thường soạn các truy vấn SQL, nhưng đôi khi bạn phải viết các truy vấn của riêng mình.

Khi bạn sử dụng Python để thực thi các truy vấn này trực tiếp vào cơ sở dữ liệu, có khả năng bạn mắc lỗi có thể ảnh hưởng đến hệ thống của bạn. Trong hướng dẫn này, bạn sẽ học cách triển khai thành công các hàm soạn các truy vấn SQL động không có khiến hệ thống của bạn gặp rủi ro khi đưa vào Python SQL.

Thiết lập cơ sở dữ liệu

Để bắt đầu, bạn sẽ thiết lập một cơ sở dữ liệu PostgreSQL mới và điền dữ liệu vào nó. Trong suốt hướng dẫn, bạn sẽ sử dụng cơ sở dữ liệu này để tận mắt chứng kiến cách thức hoạt động của Python SQL injection.

Tạo cơ sở dữ liệu

Đầu tiên, hãy mở shell của bạn và tạo một cơ sở dữ liệu PostgreSQL mới do người dùng postgres sở hữu :

$ createdb -O postgres psycopgtest

Ở đây bạn đã sử dụng tùy chọn dòng lệnh -O để đặt chủ sở hữu của cơ sở dữ liệu thành người dùng postgres . Bạn cũng đã chỉ định tên của cơ sở dữ liệu, đó là psycopgtest .

Lưu ý: postgres là người dùng đặc biệt , mà bạn thường đặt trước cho các tác vụ quản trị, nhưng đối với hướng dẫn này, bạn có thể sử dụng postgres . Tuy nhiên, trong một hệ thống thực, bạn nên tạo một người dùng riêng để làm chủ sở hữu của cơ sở dữ liệu.

Cơ sở dữ liệu mới của bạn đã sẵn sàng hoạt động! Bạn có thể kết nối với nó bằng psql :

$ psql -U postgres -d psycopgtest
psql (11.2, server 10.5)
Type "help" for help.

Bạn hiện đã kết nối với cơ sở dữ liệu psycopgtest với tư cách là người dùng postgres . Người dùng này cũng là chủ sở hữu cơ sở dữ liệu, vì vậy bạn sẽ có quyền đọc trên mọi bảng trong cơ sở dữ liệu.

Tạo bảng có dữ liệu

Tiếp theo, bạn cần tạo một bảng với một số thông tin người dùng và thêm dữ liệu vào đó:

psycopgtest=# CREATE TABLE users (
    username varchar(30),
    admin boolean
);
CREATE TABLE

psycopgtest=# INSERT INTO users
    (username, admin)
VALUES
    ('ran', true),
    ('haki', false);
INSERT 0 2

psycopgtest=# SELECT * FROM users;
 username | admin
----------+-------
 ran      | t
 haki     | f
(2 rows)

Bảng có hai cột:username và admin . admin cột cho biết người dùng có đặc quyền quản trị hay không. Mục tiêu của bạn là nhắm mục tiêu đến admin và cố gắng lạm dụng nó.

Thiết lập môi trường ảo Python

Bây giờ bạn đã có cơ sở dữ liệu, đã đến lúc thiết lập môi trường Python của bạn. Để biết hướng dẫn từng bước về cách thực hiện việc này, hãy xem Môi trường ảo Python:A Primer.

Tạo môi trường ảo của bạn trong một thư mục mới:

(~/src) $ mkdir psycopgtest
(~/src) $ cd psycopgtest
(~/src/psycopgtest) $ python3 -m venv venv

Sau khi bạn chạy lệnh này, một thư mục mới có tên venv sẽ được tạo ra. Thư mục này sẽ lưu trữ tất cả các gói bạn cài đặt bên trong môi trường ảo.

Kết nối với Cơ sở dữ liệu

Để kết nối với cơ sở dữ liệu bằng Python, bạn cần có bộ điều hợp cơ sở dữ liệu . Hầu hết các bộ điều hợp cơ sở dữ liệu tuân theo phiên bản 2.0 của Đặc tả API cơ sở dữ liệu Python PEP 249. Mỗi công cụ cơ sở dữ liệu chính đều có bộ điều hợp hàng đầu:

Cơ sở dữ liệu	Bộ điều hợp
PostgreSQL	Psycopg
SQLite	sqlite3
Oracle	cx_oracle
MySql	MySQLdb

Để kết nối với cơ sở dữ liệu PostgreSQL, bạn cần cài đặt Psycopg, đây là bộ điều hợp phổ biến nhất cho PostgreSQL bằng Python. Django ORM sử dụng nó theo mặc định và nó cũng được SQLAlchemy hỗ trợ.

Trong thiết bị đầu cuối của bạn, hãy kích hoạt môi trường ảo và sử dụng pip để cài đặt psycopg :

(~/src/psycopgtest) $ source venv/bin/activate
(~/src/psycopgtest) $ python -m pip install psycopg2>=2.8.0
Collecting psycopg2
  Using cached https://....
  psycopg2-2.8.2.tar.gz
Installing collected packages: psycopg2
  Running setup.py install for psycopg2 ... done
Successfully installed psycopg2-2.8.2

Bây giờ bạn đã sẵn sàng tạo kết nối với cơ sở dữ liệu của mình. Đây là phần bắt đầu tập lệnh Python của bạn:

import psycopg2

connection = psycopg2.connect(
    host="localhost",
    database="psycopgtest",
    user="postgres",
    password=None,
)
connection.set_session(autocommit=True)

Bạn đã sử dụng psycopg2.connect() để tạo kết nối. Hàm này chấp nhận các đối số sau:

host là địa chỉ IP hoặc DNS của máy chủ nơi đặt cơ sở dữ liệu của bạn. Trong trường hợp này, máy chủ lưu trữ là máy cục bộ của bạn hoặc localhost .
database là tên của cơ sở dữ liệu để kết nối. Bạn muốn kết nối với cơ sở dữ liệu bạn đã tạo trước đó, psycopgtest .
user là người dùng có quyền đối với cơ sở dữ liệu. Trong trường hợp này, bạn muốn kết nối với cơ sở dữ liệu với tư cách là chủ sở hữu, vì vậy bạn chuyển người dùng postgres .
password là mật khẩu cho bất kỳ ai bạn đã chỉ định trong user . Trong hầu hết các môi trường phát triển, người dùng có thể kết nối với cơ sở dữ liệu cục bộ mà không cần mật khẩu.

Sau khi thiết lập kết nối, bạn đã định cấu hình phiên với autocommit=True . Kích hoạt autocommit có nghĩa là bạn sẽ không phải quản lý các giao dịch theo cách thủ công bằng cách phát hành commit hoặc rollback . Đây là mặc định có hầu hết các ORM. Bạn cũng sử dụng hành vi này ở đây để có thể tập trung vào việc soạn các truy vấn SQL thay vì quản lý các giao dịch.

Lưu ý: Người dùng Django có thể lấy phiên bản của kết nối được ORM sử dụng từ django.db.connection :

from django.db import connection

Thực thi một truy vấn

Bây giờ bạn có kết nối với cơ sở dữ liệu, bạn đã sẵn sàng để thực thi một truy vấn:

>>>

>>> with connection.cursor() as cursor:
...     cursor.execute('SELECT COUNT(*) FROM users')
...     result = cursor.fetchone()
... print(result)
(2,)

Bạn đã sử dụng kết nối connection đối tượng để tạo cursor . Cũng giống như một tệp bằng Python, cursor được thực hiện như một trình quản lý ngữ cảnh. Khi bạn tạo ngữ cảnh, một cursor được mở để bạn sử dụng để gửi lệnh đến cơ sở dữ liệu. Khi ngữ cảnh thoát ra, cursor đóng và bạn không thể sử dụng nó nữa.

Lưu ý: Để tìm hiểu thêm về trình quản lý ngữ cảnh, hãy xem Trình quản lý ngữ cảnh Python và Tuyên bố “với”.

Khi ở bên trong ngữ cảnh, bạn đã sử dụng cursor để thực hiện một truy vấn và tìm nạp kết quả. Trong trường hợp này, bạn đã đưa ra một truy vấn để đếm các hàng trong user bàn. Để tìm nạp kết quả từ truy vấn, bạn đã thực thi cursor.fetchone() và nhận được một tuple. Vì truy vấn chỉ có thể trả về một kết quả nên bạn đã sử dụng fetchone() . Nếu truy vấn trả về nhiều hơn một kết quả, thì bạn cần phải lặp lại cursor hoặc sử dụng một trong các fetch* khác các phương pháp.

Sử dụng tham số truy vấn trong SQL

Trong phần trước, bạn đã tạo cơ sở dữ liệu, thiết lập kết nối với nó và thực hiện một truy vấn. Truy vấn bạn đã sử dụng là tĩnh . Nói cách khác, nó không có tham số . Bây giờ bạn sẽ bắt đầu sử dụng các tham số trong các truy vấn của mình.

Trước tiên, bạn sẽ triển khai một chức năng kiểm tra xem người dùng có phải là quản trị viên hay không. is_admin() chấp nhận tên người dùng và trả về trạng thái quản trị viên của người dùng đó:

# BAD EXAMPLE. DON'T DO THIS!
def is_admin(username: str) -> bool:
    with connection.cursor() as cursor:
        cursor.execute("""
            SELECT
                admin
            FROM
                users
            WHERE
                username = '%s'
        """ % username)
        result = cursor.fetchone()
    admin, = result
    return admin

Hàm này thực thi một truy vấn để tìm nạp giá trị của admin cột cho một tên người dùng nhất định. Bạn đã sử dụng fetchone() để trả về một bộ giá trị với một kết quả duy nhất. Sau đó, bạn giải nén bộ này vào biến admin . Để kiểm tra chức năng của bạn, hãy kiểm tra một số tên người dùng:

>>>

>>> is_admin('haki')
False
>>> is_admin('ran')
True

Càng xa càng tốt. Hàm trả về kết quả mong đợi cho cả hai người dùng. Nhưng những gì về người dùng không tồn tại? Hãy xem theo dấu vết Python này:

>>>

>>> is_admin('foo')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 12, in is_admin
TypeError: cannot unpack non-iterable NoneType object

Khi người dùng không tồn tại, một TypeError được nuôi dưỡng. Điều này là do .fetchone() trả về None khi không tìm thấy kết quả và giải nén None tăng TypeError . Nơi duy nhất bạn có thể giải nén một bộ tuple là nơi bạn điền admin từ result .

Để xử lý những người dùng không tồn tại, hãy tạo một trường hợp đặc biệt cho khi result là None :

# BAD EXAMPLE. DON'T DO THIS!
def is_admin(username: str) -> bool:
    with connection.cursor() as cursor:
        cursor.execute("""
            SELECT
                admin
            FROM
                users
            WHERE
                username = '%s'
        """ % username)
        result = cursor.fetchone()

    if result is None:
        # User does not exist
        return False

    admin, = result
    return admin

Tại đây, bạn đã thêm một trường hợp đặc biệt để xử lý None . Nếu username không tồn tại, khi đó hàm sẽ trả về False . Một lần nữa, hãy kiểm tra chức năng trên một số người dùng:

>>>

>>> is_admin('haki')
False
>>> is_admin('ran')
True
>>> is_admin('foo')
False

Tuyệt quá! Giờ đây, hàm cũng có thể xử lý các tên người dùng không tồn tại.

Khai thác các tham số truy vấn với Python SQL Injection

Trong ví dụ trước, bạn đã sử dụng nội suy chuỗi để tạo truy vấn. Sau đó, bạn thực hiện truy vấn và gửi chuỗi kết quả trực tiếp đến cơ sở dữ liệu. Tuy nhiên, có điều gì đó bạn có thể đã bỏ qua trong quá trình này.

Nghĩ lại username đối số bạn đã chuyển cho is_admin() . Chính xác thì biến này đại diện cho điều gì? Bạn có thể cho rằng username chỉ là một chuỗi đại diện cho tên của người dùng thực tế. Tuy nhiên, như bạn sắp thấy, kẻ xâm nhập có thể dễ dàng khai thác kiểu giám sát này và gây ra tác hại lớn bằng cách thực hiện chèn SQL trong Python.

Hãy thử kiểm tra xem người dùng sau có phải là quản trị viên hay không:

>>>

>>> is_admin("'; select true; --")
True

Chờ đã… chuyện gì vừa xảy ra vậy?

Hãy cùng nhìn lại cách triển khai. In ra truy vấn thực tế đang được thực thi trong cơ sở dữ liệu:

>>>

>>> print("select admin from users where username = '%s'" % "'; select true; --")
select admin from users where username = ''; select true; --'

Văn bản kết quả chứa ba câu lệnh. Để hiểu chính xác cách thức hoạt động của Python SQL injection, bạn cần phải kiểm tra từng phần riêng lẻ. Câu lệnh đầu tiên như sau:

select admin from users where username = '';

Đây là truy vấn dự định của bạn. Dấu chấm phẩy (; ) kết thúc truy vấn, vì vậy kết quả của truy vấn này không quan trọng. Tiếp theo là câu lệnh thứ hai:

select true;

Tuyên bố này được xây dựng bởi kẻ xâm nhập. Nó được thiết kế để luôn trả về True .

Cuối cùng, bạn thấy đoạn mã ngắn này:

--'

Đoạn mã này ngăn chặn bất cứ điều gì xuất hiện sau nó. Kẻ xâm nhập đã thêm ký hiệu nhận xét (-- ) để biến mọi thứ bạn có thể đã đặt sau trình giữ chỗ cuối cùng thành nhận xét.

Khi bạn thực thi hàm với đối số này, nó sẽ luôn trả về True . Ví dụ:nếu bạn sử dụng chức năng này trong trang đăng nhập của mình, kẻ xâm nhập có thể đăng nhập bằng tên người dùng '; select true; -- và họ sẽ được cấp quyền truy cập.

Nếu bạn nghĩ điều này là tồi tệ, nó có thể trở nên tồi tệ hơn! Những kẻ xâm nhập có kiến thức về cấu trúc bảng của bạn có thể sử dụng Python SQL injection để gây ra thiệt hại vĩnh viễn. Ví dụ:kẻ xâm nhập có thể đưa ra một câu lệnh cập nhật để thay đổi thông tin trong cơ sở dữ liệu:

>>>

>>> is_admin('haki')
False
>>> is_admin("'; update users set admin = 'true' where username = 'haki'; select true; --")
True
>>> is_admin('haki')
True

Hãy phân tích nó một lần nữa:

';

Đoạn mã này kết thúc truy vấn, giống như trong lần tiêm trước. Câu lệnh tiếp theo như sau:

update users set admin = 'true' where username = 'haki';

Phần này cập nhật admin thành true cho người dùng haki .

Cuối cùng, có đoạn mã này:

select true; --

Như trong ví dụ trước, phần này trả về true và nhận xét mọi thứ sau nó.

Tại sao điều này lại tồi tệ hơn? Chà, nếu kẻ xâm nhập quản lý để thực thi chức năng với đầu vào này, thì người dùng haki sẽ trở thành quản trị viên:

psycopgtest=# select * from users;
 username | admin
----------+-------
 ran      | t
 haki     | t
(2 rows)

Kẻ xâm nhập không còn phải sử dụng hack. Họ chỉ có thể đăng nhập bằng tên người dùng haki . (Nếu kẻ xâm nhập thực sự muốn gây hại, thì họ thậm chí có thể phát hành DROP DATABASE lệnh.)

Trước khi bạn quên, hãy khôi phục haki trở lại trạng thái ban đầu:

psycopgtest=# update users set admin = false where username = 'haki';
UPDATE 1

Vậy tại sao điều này xảy ra? Chà, bạn biết gì về username lý lẽ? Bạn biết đó phải là một chuỗi đại diện cho tên người dùng, nhưng bạn không thực sự kiểm tra hoặc thực thi xác nhận này. Điều này có thể nguy hiểm! Đó chính xác là thứ mà những kẻ tấn công đang tìm kiếm khi họ cố gắng xâm nhập hệ thống của bạn.

Tạo tham số truy vấn an toàn

Trong phần trước, bạn đã thấy cách kẻ xâm nhập có thể khai thác hệ thống của bạn và giành được quyền quản trị bằng cách sử dụng một chuỗi được chế tạo cẩn thận. Vấn đề là bạn đã cho phép giá trị được truyền từ máy khách được thực thi trực tiếp đến cơ sở dữ liệu mà không thực hiện bất kỳ loại kiểm tra hoặc xác thực nào. Việc đưa vào SQL dựa trên loại lỗ hổng này.

Bất kỳ lúc nào thông tin nhập của người dùng được sử dụng trong truy vấn cơ sở dữ liệu, sẽ có một lỗ hổng có thể xảy ra đối với việc đưa vào SQL. Chìa khóa để ngăn chặn việc tiêm SQL trong Python là đảm bảo rằng giá trị đang được sử dụng như dự định của nhà phát triển. Trong ví dụ trước, bạn dự định cho username được sử dụng như một chuỗi. Trong thực tế, nó được sử dụng như một câu lệnh SQL thô.

Để đảm bảo các giá trị được sử dụng như mong muốn, bạn cần phải thoát giá trị. Ví dụ:để ngăn những kẻ xâm nhập chèn SQL thô vào vị trí của một đối số chuỗi, bạn có thể thoát khỏi dấu ngoặc kép:

>>>

>>> # BAD EXAMPLE. DON'T DO THIS!
>>> username = username.replace("'", "''")

Đây chỉ là một ví dụ. Có rất nhiều ký tự đặc biệt và kịch bản cần nghĩ đến khi cố gắng ngăn chặn việc đưa Python SQL vào. Thật may mắn cho bạn, bộ điều hợp cơ sở dữ liệu hiện đại, đi kèm với các công cụ tích hợp để ngăn chặn việc đưa Python SQL vào bằng cách sử dụng tham số truy vấn . Chúng được sử dụng thay vì nội suy chuỗi đơn giản để soạn truy vấn với các tham số.

Lưu ý: Các bộ điều hợp, cơ sở dữ liệu và ngôn ngữ lập trình khác nhau tham chiếu đến các tham số truy vấn bằng các tên khác nhau. Các tên thông thường bao gồm biến ràng buộc , biến thay thế và biến thay thế .

Giờ bạn đã hiểu rõ hơn về lỗ hổng bảo mật, bạn đã sẵn sàng viết lại hàm bằng cách sử dụng các tham số truy vấn thay vì nội suy chuỗi:

 1def is_admin(username: str) -> bool:
 2    with connection.cursor() as cursor:
 3        cursor.execute("""
 4            SELECT
 5                admin
 6            FROM
 7                users
 8            WHERE
 9                username = %(username)s
10        """, {
11            'username': username
12        })
13        result = cursor.fetchone()
14
15    if result is None:
16        # User does not exist
17        return False
18
19    admin, = result
20    return admin

Dưới đây là những gì khác biệt trong ví dụ này:

Ở dòng 9, bạn đã sử dụng tham số có tên username để cho biết tên người dùng sẽ đi đâu. Lưu ý cách tham số username không còn được bao quanh bởi các dấu ngoặc kép.
Trong dòng 11, bạn đã chuyển giá trị của username làm đối số thứ hai cho cursor.execute() . Kết nối sẽ sử dụng loại và giá trị của username khi thực hiện truy vấn trong cơ sở dữ liệu.

Để kiểm tra chức năng này, hãy thử một số giá trị hợp lệ và không hợp lệ, bao gồm cả chuỗi nguy hiểm trước đó:

>>>

>>> is_admin('haki')
False
>>> is_admin('ran')
True
>>> is_admin('foo')
False
>>> is_admin("'; select true; --")
False

Kinh ngạc! Hàm trả về kết quả mong đợi cho tất cả các giá trị. Hơn nữa, chuỗi nguy hiểm không còn hoạt động nữa. Để hiểu lý do tại sao, bạn có thể kiểm tra truy vấn được tạo bởi execute() :

>>>

>>> with connection.cursor() as cursor:
...    cursor.execute("""
...        SELECT
...            admin
...        FROM
...            users
...        WHERE
...            username = %(username)s
...    """, {
...        'username': "'; select true; --"
...    })
...    print(cursor.query.decode('utf-8'))
SELECT
    admin
FROM
    users
WHERE
    username = '''; select true; --'

Kết nối xử lý giá trị của username dưới dạng một chuỗi và thoát khỏi bất kỳ ký tự nào có thể kết thúc chuỗi và đưa vào Python SQL.

Chuyển các tham số truy vấn an toàn

Bộ điều hợp cơ sở dữ liệu thường cung cấp một số cách để chuyển các tham số truy vấn. Trình giữ chỗ được đặt tên thường là tốt nhất để dễ đọc, nhưng một số triển khai có thể có lợi khi sử dụng các tùy chọn khác.

Hãy cùng xem nhanh một số cách sử dụng tham số truy vấn đúng và sai. Khối mã sau đây hiển thị các loại truy vấn bạn sẽ muốn tránh:

# BAD EXAMPLES. DON'T DO THIS!
cursor.execute("SELECT admin FROM users WHERE username = '" + username + '");
cursor.execute("SELECT admin FROM users WHERE username = '%s' % username);
cursor.execute("SELECT admin FROM users WHERE username = '{}'".format(username));
cursor.execute(f"SELECT admin FROM users WHERE username = '{username}'");

Mỗi câu lệnh trong số này chuyển username từ máy khách trực tiếp đến cơ sở dữ liệu mà không cần thực hiện bất kỳ loại kiểm tra hoặc xác nhận nào. Loại mã này đã chín muồi để mời đưa Python SQL vào.

Ngược lại, những loại truy vấn này sẽ an toàn để bạn thực thi:

# SAFE EXAMPLES. DO THIS!
cursor.execute("SELECT admin FROM users WHERE username = %s'", (username, ));
cursor.execute("SELECT admin FROM users WHERE username = %(username)s", {'username': username});

Trong các câu lệnh này, username được truyền dưới dạng một tham số được đặt tên. Bây giờ, cơ sở dữ liệu sẽ sử dụng loại và giá trị được chỉ định của username khi thực thi truy vấn, cung cấp khả năng bảo vệ khỏi việc tiêm SQL trong Python.

Sử dụng cấu phần SQL

Cho đến nay, bạn đã sử dụng các tham số cho các ký tự. Chữ viết là các giá trị như số, chuỗi và ngày. Nhưng nếu bạn gặp trường hợp sử dụng yêu cầu soạn một truy vấn khác — một truy vấn trong đó tham số là một thứ gì đó khác, chẳng hạn như tên bảng hoặc cột?

Lấy cảm hứng từ ví dụ trước, hãy triển khai một hàm chấp nhận tên của bảng và trả về số hàng trong bảng đó:

# BAD EXAMPLE. DON'T DO THIS!
def count_rows(table_name: str) -> int:
    with connection.cursor() as cursor:
        cursor.execute("""
            SELECT
                count(*)
            FROM
                %(table_name)s
        """, {
            'table_name': table_name,
        })
        result = cursor.fetchone()

    rowcount, = result
    return rowcount

Cố gắng thực thi chức năng trên bảng người dùng của bạn:

>>>

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 9, in count_rows
psycopg2.errors.SyntaxError: syntax error at or near "'users'"
LINE 5:                 'users'
                        ^

Lệnh không tạo được SQL. Như bạn đã thấy, bộ điều hợp cơ sở dữ liệu xử lý biến dưới dạng một chuỗi hoặc một ký tự. Tuy nhiên, tên bảng không phải là một chuỗi đơn giản. Đây là nơi xuất hiện cấu phần SQL.

Bạn đã biết không an toàn khi sử dụng phép nội suy chuỗi để soạn SQL. May mắn thay, Psycopg cung cấp một mô-đun có tên là psycopg.sql để giúp bạn soạn các truy vấn SQL một cách an toàn. Hãy viết lại hàm bằng psycopg.sql.SQL() :

from psycopg2 import sql

def count_rows(table_name: str) -> int:
    with connection.cursor() as cursor:
        stmt = sql.SQL("""
            SELECT
                count(*)
            FROM
                {table_name}
        """).format(
            table_name = sql.Identifier(table_name),
        )
        cursor.execute(stmt)
        result = cursor.fetchone()

    rowcount, = result
    return rowcount

Có hai điểm khác biệt trong cách triển khai này. Đầu tiên, bạn đã sử dụng sql.SQL() để soạn truy vấn. Sau đó, bạn đã sử dụng sql.Identifier() để chú thích giá trị đối số table_name . (Một số nhận dạng là một cột hoặc tên bảng.)

Lưu ý: Người dùng gói phổ biến django-debug-toolbar có thể gặp lỗi trong bảng điều khiển SQL cho các truy vấn được tạo bằng psycopg.sql.SQL() . Bản sửa lỗi dự kiến sẽ được phát hành trong phiên bản 2.0.

Bây giờ, hãy thử thực thi chức năng trên user bảng:

>>>

>>> count_rows('users')
2

Tuyệt quá! Tiếp theo, hãy xem điều gì sẽ xảy ra khi bảng không tồn tại:

>>>

>>> count_rows('foo')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 11, in count_rows
psycopg2.errors.UndefinedTable: relation "foo" does not exist
LINE 5:                 "foo"
                        ^

Hàm ném UndefinedTable ngoại lệ. Trong các bước sau, bạn sẽ sử dụng ngoại lệ này như một dấu hiệu cho thấy hàm của bạn an toàn trước cuộc tấn công chèn SQL trong Python.

Lưu ý: Ngoại lệ UndefinedTable đã được thêm vào psycopg2 phiên bản 2.8. Nếu bạn đang làm việc với phiên bản Psycopg cũ hơn, thì bạn sẽ gặp một ngoại lệ khác.

Để tổng hợp tất cả lại với nhau, hãy thêm tùy chọn để đếm các hàng trong bảng lên đến một giới hạn nhất định. Tính năng này có thể hữu ích cho các bảng rất lớn. Để triển khai điều này, hãy thêm LIMIT mệnh đề cho truy vấn, cùng với các tham số truy vấn cho giá trị của giới hạn:

from psycopg2 import sql

def count_rows(table_name: str, limit: int) -> int:
    with connection.cursor() as cursor:
        stmt = sql.SQL("""
            SELECT
                COUNT(*)
            FROM (
                SELECT
                    1
                FROM
                    {table_name}
                LIMIT
                    {limit}
            ) AS limit_query
        """).format(
            table_name = sql.Identifier(table_name),
            limit = sql.Literal(limit),
        )
        cursor.execute(stmt)
        result = cursor.fetchone()

    rowcount, = result
    return rowcount

Trong khối mã này, bạn đã chú thích limit sử dụng sql.Literal() . Như trong ví dụ trước, psycopg sẽ ràng buộc tất cả các tham số truy vấn dưới dạng các ký tự khi sử dụng cách tiếp cận đơn giản. Tuy nhiên, khi sử dụng sql.SQL() , bạn cần chú thích rõ ràng từng tham số bằng sql.Identifier() hoặc sql.Literal() .

Lưu ý: Thật không may, đặc tả API Python không giải quyết ràng buộc của các số nhận dạng, mà chỉ giải quyết các ký tự. Psycopg là bộ điều hợp phổ biến duy nhất bổ sung khả năng soạn SQL một cách an toàn với cả chữ và số nhận dạng. Thực tế này khiến việc chú ý kỹ hơn khi ràng buộc các số nhận dạng càng trở nên quan trọng hơn.

Thực thi hàm để đảm bảo rằng nó hoạt động:

>>>

>>> count_rows('users', 1)
1
>>> count_rows('users', 10)
2

Bây giờ bạn thấy chức năng đang hoạt động, hãy đảm bảo rằng nó cũng an toàn:

>>>

>>> count_rows("(select 1) as foo; update users set admin = true where name = 'haki'; --", 1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 18, in count_rows
psycopg2.errors.UndefinedTable: relation "(select 1) as foo; update users set admin = true where name = '" does not exist
LINE 8:                     "(select 1) as foo; update users set adm...
                            ^

Theo dõi này cho thấy rằng psycopg thoát giá trị và cơ sở dữ liệu coi nó như một tên bảng. Vì bảng có tên này không tồn tại nên một UndefinedTable ngoại lệ đã được nêu ra và bạn không bị tấn công!

Kết luận

Bạn đã triển khai thành công một hàm tạo SQL động không có đặt hệ thống của bạn vào nguy cơ bị tiêm SQL trong Python! Bạn đã sử dụng cả ký tự và số nhận dạng trong truy vấn của mình mà không ảnh hưởng đến bảo mật.

Bạn đã học:

Cái gì Chèn SQL trong Python là gì và nó có thể được khai thác như thế nào
Cách ngăn chặn việc đưa SQL vào Python sử dụng các tham số truy vấn
Cách soạn câu lệnh SQL một cách an toàn sử dụng các ký tự và số nhận dạng làm tham số

Giờ đây, bạn có thể tạo các chương trình có thể chống lại các cuộc tấn công từ bên ngoài. Đi ra ngoài và ngăn chặn tin tặc!